added pipe logic

check_app.py

@@ -37,17 +37,17 @@ class ProgressAuraFlowPipeline(DiffusionPipeline):
         self._num_inference_steps = num_inference_steps
         self._step = 0
         
-        def progress_callback(pipe, step_index, timestep, callback_kwargs):
+        def progress_callback(step_index, timestep, callback_kwargs):
             if callback and step_index % callback_steps == 0:
-                callback(step_index, timestep, callback_kwargs)
+                # Pass self (the pipeline) to the callback
+                callback(self, step_index, timestep, callback_kwargs)
             return callback_kwargs
         
         # Monkey patch the original pipeline's progress tracking
         original_step = self.original_pipeline.scheduler.step
         def wrapped_step(*args, **kwargs):
             self._step += 1
-            if callback:
-                progress_callback(self, self._step, None, {})
+            progress_callback(self._step, None, {})
             return original_step(*args, **kwargs)
         
         self.original_pipeline.scheduler.step = wrapped_step
@@ -111,7 +111,7 @@ MODEL_CONFIGS = {
     }
 }
 
-def generate_image_with_progress(pipe, prompt, num_steps, guidance_scale=None, seed=None, progress=gr.Progress(track_tqdm=True)):
+def generate_image_with_progress(model_name,pipe, prompt, num_steps, guidance_scale=None, seed=None, progress=gr.Progress(track_tqdm=True)):
     generator = None
     if seed is not None:
         generator = torch.Generator("cuda").manual_seed(seed)
@@ -162,6 +162,8 @@ def generate_image_with_progress(pipe, prompt, num_steps, guidance_scale=None, s
             prompt,
             num_inference_steps=num_steps,
             generator=generator,
+            callback=callback,
+            callback_steps=1,
         ).images[0]
 
     return image
@@ -194,7 +196,7 @@ def create_pipeline_logic(prompt_text, model_name):
         ).to("cuda")
         
     image = generate_image_with_progress(
-        pipe, prompt_text, num_steps=num_steps, guidance_scale=guidance_scale, seed=seed, progress=progress
+        model_name,pipe, prompt_text, num_steps=num_steps, guidance_scale=guidance_scale, seed=seed, progress=progress
     )
     return f"Seed: {seed}", image