Update app.py

app.py

@@ -14,36 +14,32 @@ import numpy as np
 import torch
 from diffusers import FluxImg2ImgPipeline
 
-# Global pipe variable for lazy loading
-pipe = None
-
 # Use float16 instead of bfloat16 for T4 compatibility
 dtype = torch.float16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def get_pipe():
-    global pipe
-    if pipe is None:
-        # Set more aggressive memory optimization
-        torch.cuda.empty_cache()
-        gc.collect()  # Force garbage collection
-        
-        pipe = FluxImg2ImgPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-schnell", 
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True,
-            use_safetensors=True
-        )
-        
-        # Enable attention slicing to reduce memory footprint
-        pipe.enable_attention_slicing(1)
-        
-        # Move to device more carefully
-        if torch.cuda.is_available():
-            pipe = pipe.to("cuda:0")
-        else:
-            pipe = pipe.to("cpu")
-    return pipe
+# Initialize the pipe directly during startup
+print("Loading model during startup...")
+torch.cuda.empty_cache()
+gc.collect()  # Force garbage collection
+
+pipe = FluxImg2ImgPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", 
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    use_safetensors=True
+)
+
+# Enable attention slicing to reduce memory footprint
+pipe.enable_attention_slicing(1)
+
+# Move to device immediately
+if torch.cuda.is_available():
+    pipe = pipe.to("cuda:0")
+else:
+    pipe = pipe.to("cpu")
+
+print("Model loaded successfully")
 
 def sanitize_prompt(prompt):
   # Allow only alphanumeric characters, spaces, and basic punctuation
@@ -80,12 +76,6 @@ def resize_image(image: Image.Image, max_dim: int = 384) -> Image.Image:
         image = image.resize((new_w, new_h), Image.LANCZOS)
     return image
 
-# Try to preload the model at startup
-def preload_model():
-    # Skip preloading to avoid memory issues at startup
-    print("Skipping preload - will load model on first request")
-    pass
-
 # Increase the timeout to 4 minutes
 @spaces.GPU(duration=740)
 def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=2, progress=gr.Progress(track_tqdm=True)):
@@ -94,10 +84,7 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
     torch.cuda.empty_cache()
     gc.collect()
     
-    progress(5, desc="Loading model")
-    # Get the model using lazy loading
-    model = get_pipe()
-    progress(15, desc="Model loaded")
+    progress(15, desc="Processing")
     
     def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
         if image is None:
@@ -123,7 +110,7 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
         # Use autocast for better memory efficiency
         with torch.cuda.amp.autocast(dtype=torch.float16):
             with torch.no_grad():
-                output = model(
+                output = pipe(
                     prompt=prompt, 
                     image=image, 
                     generator=generator, 
@@ -181,9 +168,6 @@ css="""
 }
 """
 
-# Try to preload the model
-preload_model()
-
 with gr.Blocks(css=css, elem_id="demo-container") as demo:
     with gr.Column():
         gr.HTML(read_file("demo_header.html"))