Update app.py

app.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 """
 RND1 Diffusion Model Demo for Hugging Face Spaces with ZeroGPU
-With Intermediate State Visualization
 """
 
 import torch
@@ -13,14 +12,12 @@ from transformers import AutoTokenizer
 from typing import Iterator
 import time
 
-# Global model and tokenizer
 model = None
 tokenizer = None
 device = "cuda"
 
 
 def set_seed(seed: int):
-    """Set random seed for reproducibility."""
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
@@ -29,7 +26,6 @@ def set_seed(seed: int):
 
 
 def load_model():
-    """Load model and tokenizer (called once at startup)."""
     global model, tokenizer
     
     from rnd.configuration_rnd import RND1Config
@@ -59,13 +55,6 @@ def load_model():
     print("Model loaded successfully!")
 
 
-def format_output_with_metadata(text: str, step_info: str = None, show_steps: bool = True) -> str:
-    """Format output with optional step information."""
-    if show_steps and step_info:
-        return f"**{step_info}**\n\n{text}"
-    return text
-
-
 @spaces.GPU(duration=120)
 def generate_with_intermediate_steps(
     prompt: str,
@@ -77,22 +66,19 @@ def generate_with_intermediate_steps(
     top_p: float,
     seed: int,
     show_intermediate: bool,
-    checkpoint_interval: int,
 ) -> Iterator[tuple[str, str]]:
-    """
-    Generate text and show intermediate states by running multiple passes.
     
-    This is a workaround for diffusion models without callback support.
-    We generate at different step counts to simulate intermediate states.
-    """
     if not prompt.strip():
-        yield "⚠️ Please enter a prompt.", "Error"
+        yield "Please enter a prompt.", "Error"
         return
     
-    yield "", "🎲 Initializing..."
+    # Keep last valid output to avoid blanks
+    last_output = ""
+    last_status = "Initializing..."
+    
+    yield last_output, last_status
     set_seed(seed)
     
-    # Format prompt based on mode
     if mode == "task":
         if not prompt.strip().startswith("Question:"):
             formatted_prompt = f"Question: {prompt}\n"
@@ -101,13 +87,12 @@ def generate_with_intermediate_steps(
     else:
         formatted_prompt = prompt
     
-    yield "", "📝 Tokenizing..."
+    last_status = "Tokenizing..."
+    yield last_output, last_status
     
-    # Tokenize
     inputs = tokenizer(formatted_prompt, return_tensors="pt")
     input_ids = inputs.input_ids.to(device)
     
-    # Prepare generation config
     from rnd.generation_config import RND1GenerationConfig
     
     greedy = (temperature == 1.0)
@@ -116,28 +101,14 @@ def generate_with_intermediate_steps(
     generator = torch.Generator(device=device)
     generator.manual_seed(seed)
     
-    if show_intermediate and num_steps >= checkpoint_interval * 2:
-        # Generate at intermediate step counts to show progression
-        # This simulates the diffusion process by stopping early
-        step_checkpoints = []
-        
-        # Create checkpoints: e.g., for 64 steps with interval 16: [16, 32, 48, 64]
-        for i in range(checkpoint_interval, num_steps + 1, checkpoint_interval):
-            step_checkpoints.append(i)
-        
-        # Make sure we always include the final step count
-        if step_checkpoints[-1] != num_steps:
-            step_checkpoints.append(num_steps)
-        
-        yield "", f"🌊 Generating with {len(step_checkpoints)} checkpoints..."
-        
-        for i, steps in enumerate(step_checkpoints):
-            # Reset generator for consistency
+    if show_intermediate:
+        # Generate at EVERY step from 1 to num_steps
+        for current_step in range(1, num_steps + 1):
             generator.manual_seed(seed)
             
             gen_config = RND1GenerationConfig(
                 max_new_tokens=max_new_tokens,
-                num_diffusion_steps=steps,
+                num_diffusion_steps=current_step,
                 mask_token_id=mask_token_id,
                 temperature=temperature if not greedy else 1.0,
                 top_k=top_k if top_k > 0 else None,
@@ -148,8 +119,7 @@ def generate_with_intermediate_steps(
                 bos_token_id=tokenizer.bos_token_id,
             )
             
-            status = f"🌊 Diffusion at {steps}/{num_steps} steps ({i+1}/{len(step_checkpoints)})"
-            yield "", status
+            last_status = f"Step {current_step}/{num_steps}"
             
             with torch.no_grad():
                 output = model.generate(
@@ -158,28 +128,21 @@ def generate_with_intermediate_steps(
                     generator=generator,
                 )
             
-            # Decode
             generated_tokens = output[0][len(input_ids[0]):]
             text = tokenizer.decode(
                 generated_tokens.tolist(),
                 skip_special_tokens=True
             )
             
-            # Show intermediate result
-            is_final = (steps == num_steps)
-            if is_final:
-                status_msg = f"✅ Final generation ({steps} steps)"
-            else:
-                status_msg = f"📊 Checkpoint {i+1}/{len(step_checkpoints)}: {steps} steps (partial quality)"
+            last_output = text
             
-            yield text, status_msg
+            if current_step == num_steps:
+                last_status = f"Complete ({num_steps} steps)"
             
-            # Small delay to make progression visible
-            if not is_final:
-                time.sleep(0.3)
+            yield last_output, last_status
     else:
-        # Single generation without intermediate steps
-        yield "", f"🚀 Generating ({num_steps} steps)..."
+        last_status = f"Generating ({num_steps} steps)..."
+        yield last_output, last_status
         
         gen_config = RND1GenerationConfig(
             max_new_tokens=max_new_tokens,
@@ -201,14 +164,14 @@ def generate_with_intermediate_steps(
                 generator=generator,
             )
         
-        # Final decode
         generated_tokens = output[0][len(input_ids[0]):]
-        final_text = tokenizer.decode(
+        last_output = tokenizer.decode(
             generated_tokens.tolist(),
             skip_special_tokens=True
         )
         
-        yield final_text, "✅ Generation complete!"
+        last_status = "Complete"
+        yield last_output, last_status
 
 
 def generate_wrapper(
@@ -221,26 +184,21 @@ def generate_wrapper(
     top_p: float,
     seed: int,
     show_intermediate: bool,
-    checkpoint_interval: int,
 ):
-    """Wrapper for Gradio to handle the generator."""
     for output, status in generate_with_intermediate_steps(
         prompt, mode, num_steps, max_new_tokens,
-        temperature, top_k, top_p, seed, show_intermediate, checkpoint_interval
+        temperature, top_k, top_p, seed, show_intermediate
     ):
         yield output, status
 
 
-# Create Gradio interface
 def create_interface():
     with gr.Blocks(title="RND1 Diffusion Language Model", theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
-        # 🌊 RND1 Diffusion Language Model
+        # RND1 Diffusion Language Model
         
-        Experience diffusion-based text generation! Unlike autoregressive models that generate left-to-right,
-        diffusion models refine **all tokens simultaneously** through iterative denoising.
-        
-        **Enable "Show Intermediate Steps"** to see how quality improves with more diffusion steps!
+        Generate text using a diffusion-based language model that refines all tokens simultaneously 
+        through iterative denoising steps.
         """)
         
         with gr.Row():
@@ -259,21 +217,11 @@ def create_interface():
                     info="Task: Q&A format | Completion: Text continuation"
                 )
                 
-                with gr.Row():
-                    show_intermediate = gr.Checkbox(
-                        label="Show Intermediate Steps",
-                        value=True,
-                        info="Display generation at different step counts to see quality improve"
-                    )
-                    
-                    checkpoint_interval = gr.Slider(
-                        minimum=8,
-                        maximum=64,
-                        value=16,
-                        step=8,
-                        label="Checkpoint Interval",
-                        info="Show results every N steps"
-                    )
+                show_intermediate = gr.Checkbox(
+                    label="Show Live Generation",
+                    value=True,
+                    info="Display output at each diffusion step (slower but shows the process)"
+                )
                 
                 with gr.Accordion("Generation Settings", open=True):
                     num_steps = gr.Slider(
@@ -281,8 +229,8 @@ def create_interface():
                         maximum=256,
                         value=64,
                         step=16,
-                        label="Final Diffusion Steps",
-                        info="More steps = better quality (will show progression if intermediate enabled)"
+                        label="Diffusion Steps",
+                        info="More steps typically improve quality"
                     )
                     
                     max_new_tokens = gr.Slider(
@@ -329,56 +277,42 @@ def create_interface():
                         label="Random Seed"
                     )
                 
-                generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-                
-                gr.Markdown("""
-                **💡 Tip:** With intermediate steps enabled, you'll see the output at different 
-                diffusion step counts (e.g., 16→32→48→64 steps). Early steps show rougher output,
-                later steps show refined, high-quality text!
-                """)
+                generate_btn = gr.Button("Generate", variant="primary", size="lg")
             
             with gr.Column(scale=1):
                 status_box = gr.Textbox(
                     label="Status",
-                    value="Ready to generate",
+                    value="Ready",
                     lines=1,
                     interactive=False
                 )
                 
                 output = gr.Textbox(
-                    label="Generated Text (Updates with Each Checkpoint)",
+                    label="Generated Text",
                     lines=18,
                     show_copy_button=True
                 )
         
         gr.Markdown("""
-        ### 🎯 How Diffusion Generation Works
-        
-        **Unlike ChatGPT-style models that generate word-by-word:**
+        ### How it works
         
-        1. **Initialize**: Create a sequence of masked/random tokens (all at once)
-        2. **Denoise**: Predict what all tokens should be simultaneously
-        3. **Refine**: Repeat the denoising process multiple times
-        4. **Result**: After N steps, you get coherent text
+        Diffusion models generate text differently than standard language models:
         
-        **With intermediate steps enabled**, you can see:
-        - **16 steps**: Rough, partially coherent text
-        - **32 steps**: Better structure and coherence
-        - **64 steps**: High-quality, well-formed output
-        - **128+ steps**: Diminishing returns, subtle improvements
+        1. Initialize all tokens as noise/masks simultaneously
+        2. Iteratively denoise and refine all tokens together
+        3. After N steps, the output converges to coherent text
         
-        ### 📊 Examples
-        Try these to see diffusion in action:
+        With live generation enabled, you can watch the text improve step by step.
         """)
         
         gr.Examples(
             examples=[
-                ["Write a Python function that finds the longest common subsequence of two strings.", "task", 64, 256, 1.0, 0, 0.0, 12345, True, 16],
-                ["Explain quantum entanglement to a 10-year-old.", "task", 64, 200, 1.0, 0, 0.0, 42, True, 16],
-                ["The most important discovery in the history of science was", "completion", 64, 256, 1.0, 0, 0.0, 9876, True, 16],
-                ["In a world where time flows backwards,", "completion", 128, 300, 1.0, 0, 0.0, 7777, True, 32],
+                ["Write a Python function that finds the longest common subsequence of two strings.", "task", 64, 256, 1.0, 0, 0.0, 12345, True],
+                ["Explain quantum entanglement to a 10-year-old.", "task", 64, 200, 1.0, 0, 0.0, 42, True],
+                ["The most important discovery in the history of science was", "completion", 64, 256, 1.0, 0, 0.0, 9876, True],
+                ["In a world where time flows backwards,", "completion", 128, 300, 1.0, 0, 0.0, 7777, False],
             ],
-            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed, show_intermediate, checkpoint_interval],
+            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed, show_intermediate],
             outputs=[output, status_box],
             fn=generate_wrapper,
             cache_examples=False,
@@ -386,7 +320,7 @@ def create_interface():
         
         generate_btn.click(
             fn=generate_wrapper,
-            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed, show_intermediate, checkpoint_interval],
+            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed, show_intermediate],
             outputs=[output, status_box],
         )
     
@@ -394,10 +328,8 @@ def create_interface():
 
 
 if __name__ == "__main__":
-    # Load model at startup
     load_model()
     
-    # Launch Gradio interface
     demo = create_interface()
-    demo.queue(max_size=10)  # Enable queue for ZeroGPU
+    demo.queue(max_size=10)
     demo.launch()