Update app.py

app.py

@@ -1,8 +1,9 @@
-# app.py
-
 import gradio as gr
 from PIL import Image
 import torch
+import os
+import json
+from datetime import datetime
 from diffusers import StableDiffusionXLImg2ImgPipeline
 from utils.planner import (
     extract_scene_plan,
@@ -11,50 +12,41 @@ from utils.planner import (
 )
 
 # ----------------------------
-# 🖥️ Device Setup
+# 💻 Device Configuration
 # ----------------------------
-use_gpu = torch.cuda.is_available()
-device = "cuda" if use_gpu else "cpu"
-dtype = torch.float16 if use_gpu else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
 # ----------------------------
-# 🧠 Load SDXL Img2Img Pipeline (no ControlNet)
+# 🧠 Load Stable Diffusion XL Img2Img Pipeline
 # ----------------------------
 pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=dtype,
-    variant="fp16" if use_gpu else None,
     use_safetensors=True,
+    variant="fp16" if device == "cuda" else None,
 )
 pipe.to(device)
-pipe.enable_model_cpu_offload()
+
+if device == "cuda":
+    pipe.enable_model_cpu_offload()
 pipe.enable_attention_slicing()
 
 # ----------------------------
-# 🎯 Image Generation Function
+# 🎨 Core Generation Function
 # ----------------------------
 def process_image(prompt, image, num_variations):
     try:
-        print("🧠 Prompt received:", prompt)
         if image is None:
             raise ValueError("🚫 Please upload an image.")
 
-        # Step 1: Analyze image + prompt → structured plan
+        print("🧠 Prompt received:", prompt)
         scene_plan = extract_scene_plan(prompt, image)
-        print("📋 Scene Plan:", scene_plan)
-
-        # Step 2: Generate prompt variations
         enriched_prompts = generate_prompt_variations_from_scene(scene_plan, prompt, num_variations)
-        print("✅ Enriched Prompts:", enriched_prompts)
-
-        # Step 3: Generate negative prompt
         negative_prompt = generate_negative_prompt_from_scene(scene_plan)
-        print("🚫 Negative Prompt:", negative_prompt)
 
-        # Step 4: Resize image to 1024x1024
         image = image.resize((1024, 1024)).convert("RGB")
 
-        # Step 5: Generate variations
         outputs = []
         for i, enriched_prompt in enumerate(enriched_prompts):
             print(f"✨ Generating Image {i + 1}...")
@@ -68,31 +60,68 @@ def process_image(prompt, image, num_variations):
             )
             outputs.append(result.images[0])
 
-        return outputs
+        # Save logs
+        log_data = {
+            "timestamp": datetime.now().isoformat(),
+            "prompt": prompt,
+            "scene_plan": scene_plan,
+            "enriched_prompts": enriched_prompts,
+            "negative_prompt": negative_prompt,
+            "device": device,
+            "num_variations": num_variations
+        }
+        os.makedirs("logs", exist_ok=True)
+        with open("logs/generation_logs.jsonl", "a") as log_file:
+            log_file.write(json.dumps(log_data) + "\n")
+
+        return outputs, "✅ Generation complete. Images ready for download."
 
     except Exception as e:
         print("❌ Generation failed:", e)
-        return [Image.new("RGB", (512, 512), color="red")]
+        return [Image.new("RGB", (512, 512), color="red")], f"❌ Error: {str(e)}"
 
 # ----------------------------
-# 🖼️ Gradio Interface
+# 🧪 Gradio Interface
 # ----------------------------
-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"),
-    title="NewCrux Product Image Generator (SDXL Only)",
-    description="Upload a product image and describe your vision. NewCrux will generate multiple enriched image variations using SDXL + AI prompt planning."
-)
+with gr.Blocks(title="NewCrux Image-to-Image Generator") as demo:
+    gr.Markdown("### 🖼️ NewCrux: Product Lifestyle Visual Generator (SDXL)\nUpload a product image and describe the scene you'd like to see.")
+
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt")
+        input_image = gr.Image(type="pil", label="Product Image")
+        num_outputs = gr.Slider(1, 5, value=3, step=1, label="Number of Variations")
+
+    generate_btn = gr.Button("Generate Image(s)")
+
+    output_gallery = gr.Gallery(label="Generated Images", show_label=True, allow_preview=True, columns=[2], height="auto")
+    output_msg = gr.Textbox(label="Status", interactive=False)
+    download_btn = gr.File(label="Download Last Image", interactive=False)
+
+    last_images = []
+
+    def generate_and_save(prompt, image, num_outputs):
+        images, message = process_image(prompt, image, num_outputs)
+        last_images.clear()
+        last_images.extend(images)
+
+        # Save the last image for download
+        file_path = "outputs/last_generated.png"
+        os.makedirs("outputs", exist_ok=True)
+        images[0].save(file_path)
+
+        return images, message, file_path
+
+    generate_btn.click(
+        generate_and_save,
+        inputs=[prompt, input_image, num_outputs],
+        outputs=[output_gallery, output_msg, download_btn]
+    )
 
 # ----------------------------
-# 🚀 Launch
+# 🚀 Launch App
 # ----------------------------
 if __name__ == "__main__":
     demo.launch()
 
 
+