Update app.py

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import torch
 import os
 import json
+import zipfile
 from datetime import datetime
 from diffusers import StableDiffusionXLImg2ImgPipeline
 from utils.planner import (
@@ -48,6 +49,10 @@ def process_image(prompt, image, num_variations):
         image = image.resize((1024, 1024)).convert("RGB")
 
         outputs = []
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        out_dir = f"outputs/session_{timestamp}"
+        os.makedirs(out_dir, exist_ok=True)
+
         for i, enriched_prompt in enumerate(enriched_prompts):
             print(f"✨ Generating Image {i + 1}...")
             result = pipe(
@@ -58,11 +63,13 @@ def process_image(prompt, image, num_variations):
                 guidance_scale=7.5,
                 num_inference_steps=30,
             )
-            outputs.append(result.images[0])
+            output_img = result.images[0]
+            output_img.save(f"{out_dir}/generated_{i+1}.png")
+            outputs.append(output_img)
 
-        # Save logs
+        # Save log
         log_data = {
-            "timestamp": datetime.now().isoformat(),
+            "timestamp": timestamp,
             "prompt": prompt,
             "scene_plan": scene_plan,
             "enriched_prompts": enriched_prompts,
@@ -74,47 +81,45 @@ def process_image(prompt, image, num_variations):
         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."
+        # Create ZIP of outputs
+        # Handle single or multiple image download
+        if num_variations == 1:
+            single_img_path = f"{out_dir}/generated_1.png"
+            return outputs, "✅ Generated one image. Ready for download.", single_img_path
+        else:
+            zip_path = f"{out_dir}/all_images.zip"
+            with zipfile.ZipFile(zip_path, "w") as zipf:
+                for i in range(len(outputs)):
+                    img_path = f"{out_dir}/generated_{i+1}.png"
+                    zipf.write(img_path, os.path.basename(img_path))
+            return outputs, f"✅ Generated {num_variations} images. Download below.", zip_path
+
 
     except Exception as e:
         print("❌ Generation failed:", e)
-        return [Image.new("RGB", (512, 512), color="red")], f"❌ Error: {str(e)}"
+        return [Image.new("RGB", (512, 512), color="red")], f"❌ Error: {str(e)}", None
 
 # ----------------------------
 # 🧪 Gradio Interface
 # ----------------------------
 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.")
+    gr.Markdown("### 🖼️ NewCrux: Product Lifestyle Visual Generator (SDXL + Prompt AI)\nUpload a product image and describe the visual you want. The system will generate realistic marketing images using AI.")
 
     with gr.Row():
-        prompt = gr.Textbox(label="Prompt")
+        prompt = gr.Textbox(label="Prompt", placeholder="e.g., A person running on the beach wearing the product")
         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)
+    generate_btn = gr.Button("🚀 Generate Image(s)")
 
-        return images, message, file_path
+    output_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=[2], height="auto")
+    output_msg = gr.Textbox(label="Generation Status", interactive=False)
+    download_zip = gr.File(label="⬇️ Download All Images (.zip)", interactive=False)
 
     generate_btn.click(
-        generate_and_save,
+        fn=process_image,
         inputs=[prompt, input_image, num_outputs],
-        outputs=[output_gallery, output_msg, download_btn]
+        outputs=[output_gallery, output_msg, download_zip]
     )
 
 # ----------------------------