Update app.py

app.py

@@ -12,6 +12,9 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+# ---------------------------------------------------------
+# 1. LOAD MODELS GLOBALLY (TO CPU/RAM ONLY)
+# ---------------------------------------------------------
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -32,11 +35,12 @@ result = ControlNetModel_Union._load_pretrained_model(
 )
 
 model = result[0]
-model = model.to(device="cuda", dtype=torch.float16)
+# ERROR FIX: Do not move to CUDA yet. Keep on CPU.
+model = model.to(dtype=torch.float16)
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+) # ERROR FIX: Removed .to("cuda")
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
@@ -44,11 +48,14 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
     vae=vae,
     controlnet=model,
     variant="fp16",
-).to("cuda")
+) # ERROR FIX: Removed .to("cuda")
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
 
+# ---------------------------------------------------------
+# HELPER FUNCTIONS
+# ---------------------------------------------------------
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     if alignment in ("Left", "Right") and source_width >= target_width:
         return False
@@ -142,8 +149,15 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
     preview = Image.alpha_composite(preview, red_mask)
     return preview
 
+# ---------------------------------------------------------
+# 2. INFERENCE (MOVE TO GPU HERE)
+# ---------------------------------------------------------
 @spaces.GPU(duration=24)
 def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+    # ERROR FIX: Move pipe to CUDA *inside* the decorated function
+    # This automatically handles the VAE and ControlNet inside the pipe
+    pipe.to("cuda")
+    
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
     
     if not can_expand(background.width, background.height, width, height, alignment):
@@ -154,7 +168,6 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
 
     final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
 
-    # --- CHANGED HERE: We run generation but DO NOT yield inside the loop ---
     with torch.autocast(device_type="cuda", dtype=torch.float16):
         (
             prompt_embeds,
@@ -163,7 +176,7 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
             negative_pooled_prompt_embeds,
         ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-        # Iterate through pipe but discard intermediate results
+        # Iterate but DO NOT YIELD yet (prevents blurry lambda response)
         for _ in pipe(
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
@@ -172,15 +185,18 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
             image=cnet_image,
             num_inference_steps=num_inference_steps
         ):
-            pass # Do nothing, just wait for it to finish
+            pass
 
-    # Only yield ONCE at the very end (The Final Image)
-    image = _ # get the last yielded image from the loop
+    # Retrieve last generated image
+    image = _
     image = image.convert("RGBA")
     cnet_image.paste(image, (0, 0), mask)
 
     yield background, cnet_image
 
+# ---------------------------------------------------------
+# UI SETUP
+# ---------------------------------------------------------
 def clear_result():
     return gr.update(value=None)