Update app.py

app.py

@@ -302,41 +302,61 @@ print("Model setup complete. Launching Gradio demo...")
 # --- Gradio Generation Function ---
 @spaces.GPU
 def generate_text(steps):
+    """
+    Generator function that yields denoised text at each step.
+    This logic is a 1:1 copy of the original Colab notebook's sampling loop.
+    """
     steps = int(steps)
     eps = 1e-5
     timesteps = torch.linspace(1, eps, steps + 1, device=DEVICE)
     step_size = (1 - eps) / steps
 
+    # Start with a fresh random sample
     x = torch.randint(0, vocab_size, (1, CONTEXT_LENGTH), device=DEVICE)
 
+    # Initial random text
     initial_text = decode(x)
     yield f"Step 0/{steps} (Initial Noise):\n\n{wrap_text(initial_text)}"
     time.sleep(0.5)
 
     with torch.no_grad():
-        for i in range(steps):
+        for i in range(steps + 1):
+
             t = timesteps[i] * torch.ones(x.shape[0], 1, device=DEVICE)
             curr_sigma_bar, _ = NOISE(t)
 
-            next_t = t - step_size
-            next_sigma_bar, _ = NOISE(next_t)
-            delta_sigma = curr_sigma_bar - next_sigma_bar
-
-            log_score = model(x, curr_sigma_bar)
-            score = torch.exp(log_score)
-
-            stag_score = staggered_score(score, delta_sigma)
-            probs = stag_score * transition(x, delta_sigma)
-            x = sample_categorical(probs)
-
+            if i < steps:
+                # This is an intermediate denoising step
+                next_sigma_bar, _ = NOISE(t - step_size)
+                delta_sigma = curr_sigma_bar - next_sigma_bar
+
+                log_score = model(x, curr_sigma_bar)
+                score = torch.exp(log_score)
+                stag_score = staggered_score(score, delta_sigma)
+                probs = stag_score * transition(x, delta_sigma)
+                x = sample_categorical(probs)
+
+            else:
+                # This is the final, full denoising step
+                # The "next sigma" is 0, so delta_sigma is the entire current noise.
+                delta_sigma = curr_sigma_bar
+
+                log_score = model(x, curr_sigma_bar)
+                score = torch.exp(log_score)
+                stag_score = staggered_score(score, delta_sigma)
+                probs = stag_score * transition(x, delta_sigma)
+                x = sample_categorical(probs)
+
+            # Yield the decoded text after each step
+            # The last yield will be the final result
             decoded_text = decode(x)
-            yield f"Step {i+1}/{steps}:\n\n{wrap_text(decoded_text)}"
-
-    final_text = decode(x)
-    yield f"Final Result (Step {steps}/{steps}):\n\n{wrap_text(final_text)}"
+            if i < steps:
+                yield f"Step {i+1}/{steps}:\n\n{wrap_text(decoded_text)}"
+            else:
+                yield f"Final Result (Step {steps}/{steps}):\n\n{wrap_text(decoded_text)}"
 
 # --- Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks(theme=gr.themes.Citrus()) as demo:
     gr.Markdown(
         """
         # The Annotated Discrete Diffusion Model: Live Demo