Update app.py

app.py

@@ -12,10 +12,8 @@ MAX_SEED = np.iinfo(np.int32).max
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-progress=gr.Progress()
-
 @spaces.GPU
-def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28):
+def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
     """
     Perform image editing using the FLUX.1 Kontext pipeline.
     
@@ -55,15 +53,6 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
         ...     guidance_scale=2.5
         ... )
     """
-
-    progress(0,desc="Starting")
-
-    def callback_fn(pipe, step, timestep, callback_kwargs):
-        print(f"[Step {step}] Timestep: {timestep}")
-        progress_value = (step+1.0)/steps
-        progress(progress_value, desc=f"Image generating, {step + 1}/{steps} steps")
-        return callback_kwargs
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
@@ -76,7 +65,6 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
             width = input_image.size[0],
             height = input_image.size[1],
             num_inference_steps=steps,
-            callback_on_step_end=callback_fn,
             generator=torch.Generator().manual_seed(seed),
         ).images[0]
     else:
@@ -84,11 +72,8 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
             prompt=prompt,
             guidance_scale=guidance_scale,
             num_inference_steps=steps,
-            callback_on_step_end=callback_fn,
             generator=torch.Generator().manual_seed(seed),
         ).images[0]
-
-    progress(1, desc="Complete")
     return image, seed, gr.Button(visible=True)
 
 @spaces.GPU