added ProgressAuraFlowPipeline

check_app.py

@@ -12,6 +12,60 @@ from diffusers import (
     LuminaText2ImgPipeline,AutoPipelineForText2Image
 )
 import gradio as gr
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+class ProgressAuraFlowPipeline(DiffusionPipeline):
+    def __init__(self, original_pipeline):
+        super().__init__()
+        self.original_pipeline = original_pipeline
+        # Register all components from the original pipeline
+        for attr_name, attr_value in vars(original_pipeline).items():
+            setattr(self, attr_name, attr_value)
+    
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt,
+        num_inference_steps=30,
+        generator=None,
+        guidance_scale=7.5,
+        callback=None,
+        callback_steps=1,
+        **kwargs
+    ):
+        # Initialize the progress tracking
+        self._num_inference_steps = num_inference_steps
+        self._step = 0
+        
+        def progress_callback(pipe, step_index, timestep, callback_kwargs):
+            if callback and step_index % callback_steps == 0:
+                callback(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, {})
+            return original_step(*args, **kwargs)
+        
+        self.original_pipeline.scheduler.step = wrapped_step
+        
+        try:
+            # Call the original pipeline
+            result = self.original_pipeline(
+                prompt=prompt,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                guidance_scale=guidance_scale,
+                **kwargs
+            )
+            
+            return result
+        finally:
+            # Restore the original step function
+            self.original_pipeline.scheduler.step = original_step
 
 cache_dir = '/workspace/hf_cache'