0

app.py

@@ -13,25 +13,23 @@ import tempfile
 # Set environment variable to reduce memory fragmentation
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
-# Check if CUDA is available, fallback to CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
+# Initialize pipeline as None - will be loaded in GPU function
+pipe = None
 
-# Load pipeline with error handling for HF Spaces
-try:
-    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-refiner-1.0",
-        torch_dtype=torch_dtype,
-        variant="fp16" if device == "cuda" else None,
-        use_safetensors=True
-    )
-    
-    # Move to device
-    pipe = pipe.to(device)
-    
-    # Enable optimizations based on available hardware
-    if device == "cuda":
-        # Use CPU offloading to reduce VRAM usage on GPU
+def load_pipeline():
+    """Load the pipeline on GPU when needed"""
+    global pipe
+    if pipe is None:
+        print("Loading pipeline...")
+        pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-refiner-1.0",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+            device_map="auto"
+        )
+        
+        # Enable memory optimizations
         pipe.enable_model_cpu_offload()
         
         # Try to enable memory efficient attention
@@ -40,13 +38,9 @@ try:
         except (ModuleNotFoundError, ImportError):
             print("xformers not available, using attention slicing")
             pipe.enable_attention_slicing()
-    else:
-        # For CPU inference, enable attention slicing
-        pipe.enable_attention_slicing()
         
-except Exception as e:
-    print(f"Error loading pipeline: {e}")
-    pipe = None
+        print("Pipeline loaded successfully!")
+    return pipe
 
 
 @spaces.GPU
@@ -60,8 +54,11 @@ def img2img(
     num_inference_steps: int = 50,
     seed: int = -1,
 ):
-    if pipe is None:
-        return None, "❌ Model failed to load. Please try again later.", None
+    # Load pipeline inside GPU context
+    try:
+        pipe = load_pipeline()
+    except Exception as e:
+        return None, f"❌ Failed to load model: {str(e)}", None
         
     try:
         # Choose image source
@@ -86,9 +83,9 @@ def img2img(
 
         # Set seed and generator
         if seed == -1:
-            generator = torch.Generator(device=device)
+            generator = torch.Generator(device="cuda")
         else:
-            generator = torch.Generator(device=device).manual_seed(seed)
+            generator = torch.Generator(device="cuda").manual_seed(seed)
 
         # Validate inputs
         if not prompt.strip():
@@ -120,7 +117,7 @@ def img2img(
             "steps": num_inference_steps,
             "width": result.width,
             "height": result.height,
-            "device": device
+            "device": "cuda"
         }
 
         # Save metadata into PNG