Update app.py

app.py

@@ -127,6 +127,11 @@ class FlowMatchingPipeline:
         self.scheduler.set_timesteps(num_inference_steps, device=self.device)
         timesteps = self.scheduler.timesteps
         
+        # Scale initial latents by scheduler's init_noise_sigma for standard diffusion
+        # Flow matching uses unscaled latents and custom ODE integration
+        if not use_flow_matching:
+            latents = latents * self.scheduler.init_noise_sigma
+        
         # Denoising loop
         for i, t in enumerate(timesteps):
             if progress_callback:
@@ -135,15 +140,19 @@ class FlowMatchingPipeline:
             # Expand latents for classifier-free guidance
             latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
             
-            # Apply shift for flow matching
+            # For standard diffusion, let scheduler handle scaling
+            # For flow matching, apply custom shift-based scaling
             if use_flow_matching and shift > 0:
                 # Compute sigma from timestep with shift
                 sigma = t.float() / 1000.0
                 sigma_shifted = (shift * sigma) / (1 + (shift - 1) * sigma)
                 
-                # Scale latent input
+                # Scale latent input for flow matching
                 scaling = torch.sqrt(1 + sigma_shifted ** 2)
                 latent_model_input = latent_model_input / scaling
+            else:
+                # For standard diffusion, scale by scheduler
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
             
             # Prepare timestep
             timestep = t.expand(latent_model_input.shape[0])