Update app.py

app.py

@@ -371,6 +371,8 @@ print(f"✓ VAE loaded (scale={VAE_SCALE})")
 
 # ============================================================================
 # EULER DISCRETE FLOW MATCHING SAMPLER
+# Training uses: x_t = (1-t)*noise + t*data, v = data - noise
+# So t=0 is noise, t=1 is data. We sample from t=0 to t=1.
 # ============================================================================
 def flux_shift(t, shift=SHIFT):
     """Flux time shift: s*t / (1 + (s-1)*t)"""
@@ -416,22 +418,23 @@ def generate(
         C = 16
         L = 128  # T5 sequence length
 
-        # Start from noise (t=1 in flow matching)
+        # Start from noise (t=0 in this convention)
         x = torch.randn(1, H_lat * W_lat, C, device=DEVICE, dtype=DTYPE, generator=generator)
         
         # Position IDs
         img_ids = TinyFluxDeep.create_img_ids(1, H_lat, W_lat, DEVICE)
         txt_ids = TinyFluxDeep.create_txt_ids(L, DEVICE)
 
-        # Timesteps: 1 -> 0 with Flux shift
-        t_linear = torch.linspace(1, 0, num_inference_steps + 1, device=DEVICE)
-        timesteps = flux_shift(t_linear, shift=SHIFT)
+        # Timesteps: 0 -> 1 (noise to data) with Flux shift
+        t_linear = torch.linspace(0, 1, num_inference_steps + 1, device=DEVICE)
+        timesteps = flux_shift(t_linear, shift=SHIFT).clamp(1e-4, 1 - 1e-4)
 
-        # Euler discrete flow matching: x_{t-dt} = x_t + v * dt
+        # Euler flow matching: x_{t+dt} = x_t + v * dt
+        # v predicts direction from noise to data
         for i in range(num_inference_steps):
             t_curr = timesteps[i]
             t_next = timesteps[i + 1]
-            dt = t_next - t_curr  # Negative since going 1->0
+            dt = t_next - t_curr  # Positive since going 0->1
 
             t_batch = t_curr.unsqueeze(0)
             guidance = torch.tensor([guidance_scale], device=DEVICE, dtype=DTYPE)
@@ -493,7 +496,6 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
-                value="cat",
                 show_label=False,
                 max_lines=2,
                 placeholder="Enter your prompt...",