Increase chunk time and log progress

app.py

@@ -579,6 +579,7 @@ def _run_diffusion_steps(
     scheduler = GENERATION_PIPE.scheduler
     with torch.no_grad():
         for i in range(start_step, end_step):
+            step_start = time.perf_counter()
             t = timesteps[i]
             latent_model_input = torch.cat([latents, condition], dim=1).to(transformer_dtype)
             timestep = t.expand(latents.shape[0])
@@ -601,10 +602,16 @@ def _run_diffusion_steps(
                 )[0]
             noise_pred = noise_uncond + GUIDANCE_SCALE * (noise_pred - noise_uncond)
             latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+            step_secs = time.perf_counter() - step_start
+            print(
+                f"[diffusion] step {i + 1}/{NUM_INFERENCE_STEPS} "
+                f"(t={float(t):.1f}, {step_secs:.2f}s)",
+                flush=True,
+            )
     return latents
 
 
-@spaces.GPU(duration=50)
+@spaces.GPU(duration=60)
 def generate_latents_setup_on_gpu(resized_image, prompt, seed, height, width):
     """Encode prompt+image, prepare initial latents and condition. NO denoising.
 
@@ -669,7 +676,7 @@ def generate_latents_setup_on_gpu(resized_image, prompt, seed, height, width):
     return state
 
 
-@spaces.GPU(duration=50)
+@spaces.GPU(duration=60)
 def generate_latents_chunk_on_gpu(state, end_step):
     """Run denoising steps from state['step_idx'] to end_step. Only transformer is moved to GPU."""
     log_cuda_mem(f"start latents chunk -> step {end_step}")