Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,6 @@
 """
 Multi-Style Image Generator with Ice Crystal Effects
-Hugging Face Spaces App
+Hugging Face Spaces App - With Diffusion Progress Streaming
 """
 
 import torch
@@ -10,6 +10,8 @@ from PIL import Image
 from pathlib import Path
 from tqdm.auto import tqdm
 import gradio as gr
+import io
+import tempfile
 
 from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
 from transformers import CLIPTextModel, CLIPTokenizer
@@ -120,7 +122,40 @@ def load_models():
         raise RuntimeError(f"Failed to load models: {e}")
 
 
-def generate_with_style(
+def decode_latents_to_image(latents_to_decode):
+    """Decode latents to PIL Image."""
+    global vae, device
+    
+    with torch.no_grad():
+        latents_scaled = 1 / 0.18215 * latents_to_decode
+        image = vae.decode(latents_scaled).sample
+    
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.cpu().permute(0, 2, 3, 1).numpy()
+    image = (image[0] * 255).astype(np.uint8)
+    return Image.fromarray(image)
+
+
+def create_gif_from_frames(frames, output_path=None, duration=200):
+    """Create an animated GIF from a list of PIL Images."""
+    if not frames:
+        return None
+    
+    if output_path is None:
+        output_path = tempfile.mktemp(suffix='.gif')
+    
+    # Save as GIF
+    frames[0].save(
+        output_path,
+        save_all=True,
+        append_images=frames[1:],
+        duration=duration,
+        loop=0
+    )
+    return output_path
+
+
+def generate_with_style_streaming(
     style_file,
     prompt,
     seed=42,
@@ -131,13 +166,20 @@ def generate_with_style(
     use_ice_crystal_guidance=False,
     ice_crystal_loss_scale=50,
     guidance_frequency=10,
-    progress=None
+    preview_frequency=5
 ):
-    """Generate an image using a style embedding with optional ice crystal guidance."""
+    """
+    Generate an image with streaming updates.
+    Yields intermediate images during generation.
+    Returns final image and GIF path at the end.
+    """
     global vae, tokenizer, text_encoder, unet, scheduler, device
     
     load_models()
     
+    # Collect frames for GIF
+    frames = []
+    
     generator = torch.Generator(device=device).manual_seed(seed)
     learned_embeds_dict = torch.load(style_file, map_location=device, weights_only=True)
 
@@ -194,10 +236,7 @@ def generate_with_style(
     scheduler.set_timesteps(num_inference_steps)
     latents = latents * scheduler.init_noise_sigma
 
-    for i, t in enumerate(tqdm(scheduler.timesteps, desc="Generating")):
-        if progress:
-            progress((i + 1) / num_inference_steps, f"Step {i + 1}/{num_inference_steps}")
-            
+    for i, t in enumerate(scheduler.timesteps):
         latent_model_input = torch.cat([latents] * 2)
         latent_model_input = scheduler.scale_model_input(latent_model_input, t)
 
@@ -231,21 +270,37 @@ def generate_with_style(
                 torch.cuda.empty_cache()
 
         latents = scheduler.step(noise_pred, t, latents).prev_sample
-
-    latents = 1 / 0.18215 * latents
-
-    with torch.no_grad():
-        image = vae.decode(latents).sample
-
-    image = (image / 2 + 0.5).clamp(0, 1)
-    image = image.cpu().permute(0, 2, 3, 1).numpy()
-    image = (image[0] * 255).astype(np.uint8)
-    image = Image.fromarray(image)
-
-    return image
+        
+        # Decode and yield intermediate preview every N steps
+        if i % preview_frequency == 0 or i == num_inference_steps - 1:
+            preview_image = decode_latents_to_image(latents)
+            frames.append(preview_image)
+            
+            # Yield progress update: (step, total, current_image, gif_path)
+            yield {
+                "step": i + 1,
+                "total": num_inference_steps,
+                "image": preview_image,
+                "gif": None  # GIF not ready yet
+            }
+
+    # Final decode
+    final_image = decode_latents_to_image(latents)
+    frames.append(final_image)
+    
+    # Create GIF from all frames
+    gif_path = create_gif_from_frames(frames, duration=300)
+    
+    # Yield final result
+    yield {
+        "step": num_inference_steps,
+        "total": num_inference_steps,
+        "image": final_image,
+        "gif": gif_path
+    }
 
 
-def generate_image(
+def generate_image_streaming(
     prompt,
     style_choice,
     custom_embedding,
@@ -253,9 +308,9 @@ def generate_image(
     guidance_scale,
     use_ice_crystal,
     ice_crystal_intensity,
-    progress=gr.Progress()
+    preview_frequency
 ):
-    """Main generation function for Gradio interface."""
+    """Streaming generation function for Gradio interface."""
     
     if custom_embedding is not None:
         style_file = custom_embedding
@@ -268,16 +323,18 @@ def generate_image(
         raise gr.Error(f"Style embedding file not found: {style_file}")
     
     try:
-        image = generate_with_style(
+        for update in generate_with_style_streaming(
             style_file=style_file,
             prompt=prompt,
             seed=int(seed),
             guidance_scale=guidance_scale,
             use_ice_crystal_guidance=use_ice_crystal,
             ice_crystal_loss_scale=ice_crystal_intensity,
-            progress=progress
-        )
-        return image
+            preview_frequency=int(preview_frequency)
+        ):
+            status = f"Step {update['step']}/{update['total']}"
+            yield update["image"], update["gif"], status
+            
     except Exception as e:
         raise gr.Error(f"Generation failed: {str(e)}")
 
@@ -294,12 +351,13 @@ with gr.Blocks(
     # Multi-Style Image Generator with Ice Crystal Effects
     
     Generate images using textual inversion style embeddings with optional ice crystal overlay effects.
+    **Now with live diffusion progress streaming!**
     
     **Instructions:**
     1. Enter a prompt using `<style>` as placeholder (e.g., "A cat in the style of <style>")
     2. Select a predefined style OR upload your own `.bin` embedding file
     3. Optionally enable ice crystal effect for a crystalline overlay
-    4. Click Generate!
+    4. Click Generate and watch the image evolve!
     """)
     
     with gr.Row():
@@ -353,27 +411,42 @@ with gr.Blocks(
                     info="Higher = stronger crystal effect"
                 )
             
+            with gr.Accordion("Streaming Settings", open=True):
+                preview_frequency = gr.Slider(
+                    label="Preview Frequency",
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                    info="Show preview every N steps (lower = more updates, slower)"
+                )
+            
             generate_btn = gr.Button("Generate", variant="primary", size="lg")
+            status_text = gr.Textbox(label="Status", interactive=False, value="Ready")
         
         with gr.Column(scale=1):
             output_image = gr.Image(
-                label="Generated Image",
+                label="Live Preview / Final Image",
                 type="pil"
             )
+            output_gif = gr.File(
+                label="Diffusion Progress GIF (available after generation)",
+                type="filepath"
+            )
     
     gr.Examples(
         examples=[
-            ["A cat in the style of <style>", "8bit", None, 42, 7.5, False, 50],
-            ["A mystical forest in the style of <style>", "dr_strange", None, 123, 7.5, False, 50],
-            ["A portrait in the style of <style>", "max_naylor", None, 456, 7.5, True, 60],
+            ["A cat in the style of <style>", "8bit", None, 42, 7.5, False, 50, 5],
+            ["A mystical forest in the style of <style>", "dr_strange", None, 123, 7.5, False, 50, 5],
+            ["A portrait in the style of <style>", "max_naylor", None, 456, 7.5, True, 60, 5],
         ],
-        inputs=[prompt, style_choice, custom_embedding, seed, guidance_scale, use_ice_crystal, ice_crystal_intensity],
+        inputs=[prompt, style_choice, custom_embedding, seed, guidance_scale, use_ice_crystal, ice_crystal_intensity, preview_frequency],
     )
     
     generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, style_choice, custom_embedding, seed, guidance_scale, use_ice_crystal, ice_crystal_intensity],
-        outputs=output_image
+        fn=generate_image_streaming,
+        inputs=[prompt, style_choice, custom_embedding, seed, guidance_scale, use_ice_crystal, ice_crystal_intensity, preview_frequency],
+        outputs=[output_image, output_gif, status_text]
     )
 
 if __name__ == "__main__":