handling pipeline bug

app.py

@@ -14,21 +14,28 @@ torch_dtype = torch.bfloat16
 
 print("Starting Flux2 Image Generator...")
 
-# Load the pipeline at startup (NOT inside GPU decorator)
+# Load the pipeline at startup
 print("Loading Flux2 pipeline...")
 pipe = None
 
-try:
-    pipe = Flux2Pipeline.from_pretrained(
-        repo_id,
-        text_encoder=None,
-        torch_dtype=torch_dtype,
-        device_map="balanced"  # Use balanced for CPU during startup
-    )
-    print("Pipeline loaded successfully!")
-except Exception as e:
-    print(f"Error loading pipeline: {e}")
-    # Don't raise - will try to load later if needed
+def load_pipeline_startup():
+    """Load pipeline at startup without CUDA."""
+    global pipe
+    try:
+        print("Loading pipeline components...")
+        pipe = Flux2Pipeline.from_pretrained(
+            repo_id,
+            text_encoder=None,
+            torch_dtype=torch_dtype,
+        )
+        # Keep on CPU initially - will move to CUDA when needed
+        print("Pipeline loaded successfully on CPU!")
+    except Exception as e:
+        print(f"Warning: Could not load pipeline at startup: {e}")
+        print("Pipeline will be loaded on first use.")
+
+# Try to load at startup
+load_pipeline_startup()
 
 def remote_text_encoder(prompts):
     """Encode prompts using remote text encoder API."""
@@ -145,12 +152,16 @@ def generate_image(
     progress(0, desc="Moving model to GPU...")
     
     try:
-        # Move pipeline to GPU
+        # Load or get pipeline
         if pipe is None:
-            raise gr.Error("Pipeline not loaded. Please refresh the page.")
+            print("Pipeline not loaded at startup, loading now...")
+            load_pipeline_startup()
+            if pipe is None:
+                raise gr.Error("Failed to load pipeline. Please try again or contact support.")
         
         print("Moving pipeline to CUDA...")
-        pipe = pipe.to("cuda")
+        pipeline = pipe.to("cuda")
+        torch.cuda.empty_cache()  # Clear cache before generation
         
         progress(0.1, desc="Encoding prompt...")
         print("Encoding prompt...")
@@ -194,12 +205,17 @@ def generate_image(
         
         # Generate image
         with torch.inference_mode():
-            result = pipe(**pipe_kwargs)
+            result = pipeline(**pipe_kwargs)
             image = result.images[0]
         
         print("Generation complete!")
         progress(1.0, desc="Done!")
         
+        # Move pipeline back to CPU to free GPU memory
+        print("Moving pipeline back to CPU...")
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
+        
         return image
     
     except gr.Error: