Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
-import numpy as np
+from typing import List
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id = "nitrosocke/Ghibli-Diffusion"
@@ -10,44 +10,104 @@ model_id = "nitrosocke/Ghibli-Diffusion"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     model_id,
     torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-)
-pipe.to(device)
+    safety_checker=None
+).to(device)
 pipe.enable_attention_slicing()
 
-def generate_ghibli_style_realtime(image, steps=25):
-    prompt = "ghibli style portrait"
+styles = {
+    "Classic Ghibli": "ghibli style portrait",
+    "Spirited Forest": "studio ghibli mystical forest portrait, soft lighting",
+    "Windy Valley": "ghibli style sky valley portrait, dreamy atmosphere",
+    "Cozy Home": "ghibli style cozy cottage scene, warm tones"
+}
 
-    def stream():
-        def callback(step: int, timestep: int, latents):
-            with torch.no_grad():
-                img = pipe.decode_latents(latents)
-                img = pipe.numpy_to_pil(img)[0]
-            yield img  # Stream each image!
+def generate_final_image(
+    image: Image.Image,
+    style_choice: str,
+    steps: int,
+    history: List[Image.Image],
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> tuple:
+    prompt = styles.get(style_choice, "ghibli style portrait")
+    
+    with torch.inference_mode():
+        output = pipe(
+            prompt=prompt,
+            image=image,
+            strength=0.6,
+            guidance_scale=6.5,
+            num_inference_steps=steps
+        )
+    
+    final_image = output.images[0]
+    return final_image, history + [final_image]
 
-        # Capture all images and stream via yield
-        with torch.inference_mode():
-            pipe(
-                prompt=prompt,
-                image=image,
-                strength=0.6,
-                guidance_scale=6.0,
-                num_inference_steps=steps,
-                callback=callback,
-                callback_steps=1,
+# Rest of the Gradio interface code remains the same as previous version
+with gr.Blocks(theme=gr.themes.Soft()) as iface:
+    gr.Markdown("## 🌸 **Ghibli Portrait Generator — Spicy Edition** 🌸")
+    gr.Markdown("Upload a photo and transform it into an anime scene in your favorite Ghibli style! 🎬✨")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="📸 Upload your photo", height=300)
+            style_dropdown = gr.Dropdown(
+                list(styles.keys()),
+                label="🎨 Choose a Ghibli Style",
+                value="Classic Ghibli"
             )
+            steps_slider = gr.Slider(10, 100, value=100, step=1, label="✨ Inference Steps")
+            generate_btn = gr.Button("💫 Start Magic!", variant="primary")
+        
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="🌟 Final Ghibli Portrait", height=500)
+            with gr.Row():
+                download_btn = gr.Button("💾 Download Final Image")
+                clear_btn = gr.Button("🧹 Clear History")
+    
+    with gr.Accordion("📜 Previous Generations", open=False):
+        history_gallery = gr.Gallery(
+            label="Your Ghibli Journey",
+            columns=4,
+            height="auto",
+            object_fit="contain"
+        )
+    
+    history_state = gr.State([])
 
-    return stream()
+    # Generation workflow
+    generate_btn.click(
+        generate_final_image,
+        [image_input, style_dropdown, steps_slider, history_state],
+        [output_image, history_state]
+    ).then(
+        lambda x: x[-4:],  # Update gallery after generation
+        history_state,
+        history_gallery
+    )
+    
+    # Image selection from history
+    history_gallery.select(
+        lambda evt: evt,
+        None,
+        image_input
+    )
+    
+    # Download handler
+    download_btn.click(
+        lambda img: img,
+        output_image,
+        gr.File(label="⬇️ Your Ghibli Portrait")
+    )
+    
+    # Clear history
+    clear_btn.click(
+        lambda: [],
+        None,
+        history_state
+    ).then(
+        lambda: None,
+        None,
+        history_gallery
+    )
 
-iface = gr.Interface(
-    fn=generate_ghibli_style_realtime,
-    inputs=[
-        gr.Image(type="pil", label="Upload a photo"),
-        gr.Slider(minimum=10, maximum=50, value=25, step=1, label="Inference Steps")
-    ],
-    outputs=gr.Image(label="Live Ghibli Transformation"),
-    live=True,
-    title="✨ Real-Time Ghibli Portrait ✨",
-    description="Upload a photo and see the Ghibli transformation in real-time!"
-)
-
-iface.launch()
+iface.launch(share=True, debug=True