Update app.py

app.py

@@ -16,27 +16,21 @@ access_token = os.getenv("HF_ACCESS_TOKEN")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-# Load the Stable Diffusion model
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-3-medium",
-        torch_dtype=torch.float16,
-        use_auth_token=access_token  # Use the token here
-    )
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else:
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-3-medium",
-        use_auth_token=access_token  # Use the token here
-    )
-    pipe = pipe.to(device)
-
-logging.info("Loading the model...")
-
-# Load model
-logging.info("Model loaded successfully.")
+# Global variable for the pipeline
+pipe = None
+
+def load_model():
+    global pipe
+    if pipe is None:
+        logging.info("Loading the Stable Diffusion model...")
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-3-medium",
+            torch_dtype=torch.float16,
+            use_auth_token=access_token,
+            cache_dir="/path/to/cache"  # specify cache directory
+        )
+        pipe = pipe.to(device)
+        logging.info("Model loaded successfully.")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024