diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index 6fe2a8024..aaddfa0f7 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -89,6 +89,7 @@ def __init__(
         graph_attr: Optional[dict] = None,
         node_attr: Optional[dict] = None,
         edge_attr: Optional[dict] = None,
+        save_to_disk: bool = True,  # Add flag to control saving behavior
     ):
         """Diagram represents a global diagrams context.
 
@@ -118,6 +119,9 @@ def __init__(
             filename = "_".join(self.name.split()).lower()
         self.filename = filename
         self.dot = Digraph(self.name, filename=self.filename, strict=strict)
+        if save_to_disk is None:
+            save_to_disk = True
+
 
         # Set attributes.
         for k, v in self._default_graph_attrs.items():
@@ -152,6 +156,7 @@ def __init__(
 
         self.show = show
         self.autolabel = autolabel
+        self.save_to_disk = save_to_disk
 
     def __str__(self) -> str:
         return str(self.dot)
@@ -161,11 +166,13 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.render()
-        # Remove the graphviz file leaving only the image.
-        os.remove(self.filename)
+        if self.save_to_disk:  # Only render if save_to_disk is True
+            self.render()
+            # Remove the graphviz file leaving only the image.
+            os.remove(self.filename)
         setdiagram(None)
 
+
     def _repr_png_(self):
         return self.dot.pipe(format="png")