Update app.py

app.py

@@ -6,46 +6,46 @@ import os
 import torch
 import numpy as np
 import cv2
-from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel
+
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 from utils.planner import extract_scene_plan, generate_prompt_variations_from_scene  # 🧠 Brain Layer
 
 # ----------------------------
 # 🔧 Device Setup
 # ----------------------------
-device = "cpu"  # ✅ Force CPU since you're getting NVIDIA driver error
-dtype = torch.float32  # ✅ Avoid float16 on CPU
+device = "cpu"  # ✅ Using CPU for now
+dtype = torch.float32
 
 # ----------------------------
-# ✅ Use Valid Public Models
+# ✅ Load ControlNet + SDXL Model (Corrected)
 # ----------------------------
-# Valid ControlNet Canny model for SDXL
 controlnet = ControlNetModel.from_pretrained(
-    "diffusers/controlnet-canny-sdxl-1.0",  # ⚠️ Replace with actual model or use SD 1.5 version for now
+    "diffusers/controlnet-canny-sdxl-1.0",  # Use official Canny + SDXL ControlNet model
     torch_dtype=dtype
 )
 
-# Base SDXL model
-pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet,
     torch_dtype=dtype
 ).to(device)
 
-# ✅ Comment out GPU-only features
-# pipe.enable_xformers_memory_efficient_attention()  # Only for CUDA
-# pipe.enable_model_cpu_offload()  # Not needed on CPU
-
 # ----------------------------
-# 🖼 Canny Edge Generator
+# 🔍 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)
     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")
 
 # ----------------------------
@@ -54,35 +54,39 @@ def generate_canny_map(image: Image.Image) -> Image.Image:
 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: Extract scene plan
         scene_plan = extract_scene_plan(prompt)
         print("🧠 Scene plan extracted:", scene_plan)
 
+        # Step 2: Generate enriched prompt variations
         prompt_list = generate_prompt_variations_from_scene(scene_plan, prompt, num_variations)
         print("🧠 Enriched Prompts:")
         for i, p in enumerate(prompt_list):
             print(f"  {i+1}: {p}")
 
+        # Step 3: Prepare image and Canny edge
         image = image.resize((1024, 1024)).convert("RGB")
         canny_map = generate_canny_map(image)
 
         outputs = []
         for i, enriched_prompt in enumerate(prompt_list):
-            print(f"🎨 Generating image {i+1} with enriched prompt")
+            print(f"🎨 Generating image {i+1}...")
             try:
                 result = pipe(
                     prompt=enriched_prompt,
                     image=image,
-                    controlnet_conditioning_image=canny_map,
+                    control_image=canny_map,
                     num_inference_steps=30,
                     strength=0.5,
                     guidance_scale=7.5
                 )
                 outputs.append(result.images[0])
-            except Exception as inner:
-                print(f"❌ Failed to generate image {i+1}:", inner)
+            except Exception as err:
+                print(f"❌ Failed to generate image {i+1}:", err)
                 outputs.append(Image.new("RGB", (512, 512), color="red"))
 
         return outputs, scene_plan, canny_map
@@ -95,7 +99,7 @@ def process_image(prompt, image, num_variations):
 # 🖼 Gradio UI
 # ----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 NewCrux AI — SDXL + Canny (CPU Mode)\nUpload a product image, enter a prompt, and generate scenes using ControlNet.")
+    gr.Markdown("## 🧠 NewCrux AI — SDXL + Canny (CPU Mode)\nUpload a product image, enter a prompt, and generate enhanced visuals using ControlNet.")
 
     with gr.Row():
         with gr.Column():
@@ -117,4 +121,3 @@ with gr.Blocks() as demo:
 
 demo.launch()
 
-