Update app.py

app.py

@@ -24,34 +24,46 @@ controlnet = ControlNetModel.from_pretrained(
 )
 
 pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",  # use SDXL if desired: "stabilityai/stable-diffusion-xl-base-1.0"
+    "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet,
-    torch_dtype=dtype
+    torch_dtype=dtype,
+    variant="fp16" if dtype == torch.float16 else None
 ).to(device)
 
-pipe.enable_xformers_memory_efficient_attention() if device == "cuda" else None
+if device == "cuda":
+    pipe.enable_xformers_memory_efficient_attention()
+else:
+    pipe.enable_model_cpu_offload()
 
 # ----------------------------
 # 🖼 Canny Edge Generator
 # ----------------------------
 def generate_canny_map(image: Image.Image) -> Image.Image:
     print("🔍 Generating Canny map...")
+
+    if image is None:
+        raise ValueError("🚫 No image passed to Canny generator")
+
     image = image.resize((1024, 1024)).convert("RGB")
     np_image = np.array(image)
-    np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(np_image, 100, 200)
-    canny_image = Image.fromarray(edges).convert("RGB")
-    return canny_image
+    gray = cv2.cvtColor(np_image, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, 100, 200)
 
+    if edges is None:
+        raise ValueError("🚫 OpenCV Canny failed to produce edge map")
+
+    return Image.fromarray(edges).convert("RGB")
 
 # ----------------------------
 # 🎨 Image Generation Function
 # ----------------------------
-
 def process_image(prompt, image, num_variations):
     try:
         print("🧠 Prompt received:", prompt)
 
+        if image is None:
+            raise ValueError("🚫 Uploaded image is missing or invalid.")
+
         # Step 1: Brain Layer (Scene Plan)
         scene_plan = extract_scene_plan(prompt)
         print("🧠 Scene plan extracted:", scene_plan)
@@ -67,6 +79,9 @@ def process_image(prompt, image, num_variations):
         image = image.resize((1024, 1024)).convert("RGB")
         canny_map = generate_canny_map(image)
 
+        if canny_map is None:
+            raise ValueError("🚫 Canny map generation failed.")
+
         # Step 4: Generate images
         outputs = []
         for i, enriched_prompt in enumerate(prompt_list):
@@ -75,7 +90,7 @@ def process_image(prompt, image, num_variations):
                 result = pipe(
                     prompt=enriched_prompt,
                     image=image,
-                    controlnet_conditioning_image=canny_map,  # ✅ key fix here!
+                    controlnet_conditioning_image=canny_map,
                     num_inference_steps=40,
                     strength=0.5,
                     guidance_scale=7.5
@@ -96,8 +111,6 @@ def process_image(prompt, image, num_variations):
         print("❌ Generation failed:", e)
         return ["❌ Error during generation"], {"error": str(e)}, None
 
-
-
 # ----------------------------
 # 🖼 Gradio UI
 # ----------------------------
@@ -129,4 +142,3 @@ with gr.Blocks() as demo:
 
 demo.launch()
 
-