Update app.py

app.py

@@ -2,135 +2,92 @@
 
 import gradio as gr
 from PIL import Image
-import os
 import torch
-import numpy as np
-import cv2
 
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from diffusers import StableDiffusionXLImg2ImgPipeline
 from utils.planner import (
     extract_scene_plan,
     generate_prompt_variations_from_scene,
-    generate_negative_prompt_from_scene,
-    save_generation_log
+    generate_negative_prompt_from_scene
 )
 
 # ----------------------------
 # 🔧 Device Setup
 # ----------------------------
-device = "cpu"
-dtype = torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
 # ----------------------------
-# ✅ Load ControlNet Canny + SDXL Model
+# ✅ Load SDXL Only Pipeline
 # ----------------------------
-controlnet = ControlNetModel.from_pretrained(
-    "diffusers/controlnet-canny-sdxl-1.0",
-    torch_dtype=dtype
-)
-
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    controlnet=controlnet,
-    torch_dtype=dtype
-).to(device)
-
-# ----------------------------
-# 🔍 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")
+    torch_dtype=dtype,
+    variant="fp16" if device == "cuda" else None,
+    use_safetensors=True,
+)
+pipe.to(device)
+pipe.enable_model_cpu_offload()
+pipe.enable_attention_slicing()
 
 # ----------------------------
 # 🎨 Image Generation Function
 # ----------------------------
 def process_image(prompt, image, num_variations):
     try:
-        print("🧠 Prompt received:", prompt)
+        print("🧠 User Prompt:", prompt)
         if image is None:
-            raise ValueError("🚫 Uploaded image is missing or invalid.")
+            raise ValueError("🚫 Uploaded image is missing.")
 
-        # Step 1: Scene Planning
+        # Step 1: Extract scene plan
         scene_plan = extract_scene_plan(prompt, image)
-        print("🧠 Scene plan extracted:", scene_plan)
+        print("📋 Scene Plan:", scene_plan)
 
-        # Step 2: Prompt Variations
+        # Step 2: Generate enriched prompts
         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}")
+        print("✅ Enriched Prompts:", prompt_list)
 
-        # Step 3: Negative Prompt
+        # Step 3: Generate negative prompt
         negative_prompt = generate_negative_prompt_from_scene(scene_plan)
         print("🚫 Negative Prompt:", negative_prompt)
 
-        # ✅ Save log
-        caption = scene_plan.get("caption", "N/A")
-        save_generation_log(caption, scene_plan, prompt_list, negative_prompt)
-
-        # Step 4: Canny Edge Map
+        # Step 4: Resize image to SDXL resolution
         image = image.resize((1024, 1024)).convert("RGB")
-        canny_map = generate_canny_map(image)
 
-        # Step 5: Generate Images
+        # Step 5: Generate outputs with SDXL only
         outputs = []
         for i, enriched_prompt in enumerate(prompt_list):
-            print(f"🎨 Generating image {i+1}...")
-            try:
-                result = pipe(
-                    prompt=enriched_prompt,
-                    negative_prompt=negative_prompt,
-                    image=image,
-                    control_image=canny_map,
-                    num_inference_steps=30,
-                    strength=0.2,
-                    guidance_scale=7.5
-                )
-                outputs.append(result.images[0])
-            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
+            print(f"🎨 Generating variation {i+1}...")
+            result = pipe(
+                prompt=enriched_prompt,
+                negative_prompt=negative_prompt,
+                image=image,
+                strength=0.7,                # ← You can fine-tune this
+                guidance_scale=7.5,
+                num_inference_steps=30,
+            )
+            outputs.append(result.images[0])
+
+        return outputs
 
     except Exception as e:
-        print("❌ Generation failed:", e)
-        return ["❌ Error during generation"], {"error": str(e)}, None
+        print("❌ Generation Error:", e)
+        return [Image.new("RGB", (512, 512), color="red")]
 
 # ----------------------------
-# 🖼 Gradio UI
+# 🖼️ Gradio Interface
 # ----------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 NewCrux AI — SDXL + ControlNet Canny (CPU Mode)\nUpload a product image, enter a prompt, and generate lifestyle visuals guided by edges.")
-
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(label="Prompt")
-            image_input = gr.Image(type="pil", label="Upload Product Image")
-            variation_slider = gr.Slider(1, 4, step=1, value=1, label="Number of Variations")
-            generate_btn = gr.Button("Generate")
-
-        with gr.Column():
-            output_gallery = gr.Gallery(label="Generated Variations", columns=2, rows=2, height="auto")
-            json_output = gr.JSON(label="🧠 Brain Layer Reasoning")
-            canny_preview = gr.Image(label="🔍 Canny Edge Preview")
-
-    generate_btn.click(
-        fn=process_image,
-        inputs=[prompt_input, image_input, variation_slider],
-        outputs=[output_gallery, json_output, canny_preview]
-    )
+demo = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Image(type="pil", label="Product Image"),
+        gr.Slider(1, 5, value=3, step=1, label="Number of Variations")
+    ],
+    outputs=gr.Gallery(label="Generated Images").style(grid=[2], height="auto"),
+    title="NewCrux Product Image Generator (SDXL Only)",
+    description="Upload a product image and enter a prompt. SDXL will generate enriched variations using AI."
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()