Update app.py

app.py

@@ -80,7 +80,6 @@ pipe.to(device)
 # flash-attn估计库估计更新了，导致冲突了，不使用预编译的了
 # spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
 
-
 def image_to_data_uri(img):
     buffered = io.BytesIO()
     img.save(buffered, format="PNG")
@@ -162,27 +161,32 @@ def get_duration(prompt_embeds, image_list, width, height, num_inference_steps,
     return max(65, num_inference_steps * step_duration + 10)
 
 @spaces.GPU(duration=get_duration)
-def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress()):
     # Move embeddings to GPU only when inside the GPU decorated function
     prompt_embeds = prompt_embeds.to(device)
     
     generator = torch.Generator(device=device).manual_seed(seed)
-    
-    pipe_kwargs = {
-        "prompt_embeds": prompt_embeds,
-        "image": image_list,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-        "width": width,
-        "height": height,
-    }
-    
+ 
     # Progress bar for the actual generation steps
     if progress:
         progress(0, desc="Starting generation...")
+
+    def callback_fn(pipe, step, timestep, callback_kwargs):
+        print(f"[Step {step}] Timestep: {timestep}")
+        progress_value = (step+1.0)/num_inference_steps
+        progress(progress_value, desc=f"Image generating, {step + 1}/{num_inference_steps} steps")
+        return callback_kwargs
         
-    image = pipe(**pipe_kwargs).images[0]
+    image = pipe(
+        prompt_embeds=prompt_embeds,
+        image=image_list,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        width=width,
+        height=height,
+        callback_on_step_end=callback_fn,
+    ).images[0]
     return image
 
 def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):