Update app.py

app.py

@@ -9,29 +9,32 @@ from PIL import Image
 # Import the pipeline from diffusers
 from diffusers import FluxKontextPipeline
 
-# --- Constants and Model Loading ---
+# --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
 
-# Load the pretrained model
-try:
-    pipe = FluxKontextPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-Kontext-dev", 
-        torch_dtype=torch.bfloat16,
-    ).to("cuda")
-except Exception as e:
-    pipe = None
-    print(f"Warning: Could not load the model on CUDA. GPU is required. Error: {e}")
+# --- Global pipeline variable ---
+pipe = None
 
-# --- Core Inference Function for ChatInterface ---
+def load_model():
+    """Load the model on CPU first, then move to GPU when needed"""
+    global pipe
+    if pipe is None:
+        pipe = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", 
+            torch_dtype=torch.bfloat16,
+        )
+    return pipe
 
-@spaces.GPU
+# --- Core Inference Function for ChatInterface ---
+@spaces.GPU(duration=120)  # Set duration based on expected inference time
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
     """
     Performs image generation or editing based on user input from the chat interface.
     """
-    if pipe is None:
-        raise gr.Error("Model could not be loaded. This could be due to no access to the model or no CUDA-enabled GPU.")
-
+    # Load and move model to GPU within the decorated function
+    pipe = load_model()
+    pipe = pipe.to("cuda")
+    
     prompt = message["text"]
     files = message["files"]
 
@@ -62,6 +65,10 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
             num_inference_steps=steps,
             generator=generator,
         ).images[0]
+    
+    # Move model back to CPU to free GPU memory
+    pipe = pipe.to("cpu")
+    torch.cuda.empty_cache()
         
     return image
 
@@ -72,8 +79,7 @@ randomize_checkbox = gr.Checkbox(label="Randomize seed", value=False)
 guidance_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
 steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1)
 
-# --- FIX 2: Remove examples with external URLs that cause 403 errors ---
-# Instead, provide text-only examples that work without external image dependencies
+# --- Examples without external URLs ---
 examples = [
     [
         {"text": "A cute robot reading a book in a cozy library", "files": []}, 
@@ -93,7 +99,7 @@ demo = gr.ChatInterface(
     fn=chat_fn,
     title="FLUX.1 Kontext [dev]",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext</b> model.
+    A simple chat UI for the <b>FLUX.1 Kontext</b> model running on ZeroGPU.
     <br>
     To edit an image, upload it and type your instructions (e.g., "Add a hat").
     <br>