import spaces import torch from inspect import signature from diffusers import ( FluxPipeline, StableDiffusion3Pipeline, PixArtSigmaPipeline, SanaPipeline, AuraFlowPipeline, Kandinsky3Pipeline, HunyuanDiTPipeline, 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(step_index, timestep, callback_kwargs): if callback and step_index % callback_steps == 0: # 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 progress_callback(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' MODEL_CONFIGS = { "FLUX": { "repo_id": "black-forest-labs/FLUX.1-dev", "pipeline_class": FluxPipeline, }, "Stable Diffusion 3.5": { "repo_id": "stabilityai/stable-diffusion-3.5-large", "pipeline_class": StableDiffusion3Pipeline, }, "PixArt": { "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", "pipeline_class": PixArtSigmaPipeline, }, "SANA": { "repo_id": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", "pipeline_class": SanaPipeline, }, "AuraFlow": { "repo_id": "fal/AuraFlow", "pipeline_class": AuraFlowPipeline, }, "Kandinsky": { "repo_id": "kandinsky-community/kandinsky-3", "pipeline_class": Kandinsky3Pipeline, }, "Hunyuan": { "repo_id": "Tencent-Hunyuan/HunyuanDiT-Diffusers", "pipeline_class": HunyuanDiTPipeline, }, "Lumina": { "repo_id": "Alpha-VLLM/Lumina-Next-SFT-diffusers", "pipeline_class": LuminaText2ImgPipeline, } } 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) def callback(pipe, step_index, timestep, callback_kwargs): print(f" callback => {step_index}, {timestep}") if step_index is None: step_index = 0 cur_prg = step_index / num_steps progress(cur_prg, desc=f"Step {step_index}/{num_steps}") return callback_kwargs print(f"START GENR ") # Get the signature of the pipe pipe_signature = signature(pipe) # Check for the presence of "guidance_scale" and "callback_on_step_end" in the signature has_guidance_scale = "guidance_scale" in pipe_signature.parameters has_callback_on_step_end = "callback_on_step_end" in pipe_signature.parameters if has_guidance_scale and has_callback_on_step_end: print("has callback_on_step_end and has guidance_scale") image = pipe( prompt, num_inference_steps=num_steps, generator=generator, guidance_scale=guidance_scale, callback_on_step_end=callback, ).images[0] elif not has_callback_on_step_end and has_guidance_scale: print("NO callback_on_step_end and has guidance_scale") image = pipe( prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator, callback=callback, callback_steps=1, ).images[0] elif has_callback_on_step_end and not has_guidance_scale: print("has callback_on_step_end and NO guidance_scale") image = pipe( prompt, num_inference_steps=num_steps, generator=generator, callback_on_step_end=callback, ).images[0] elif not has_callback_on_step_end and not has_guidance_scale: print("NO callback_on_step_end and NO guidance_scale") image = pipe( prompt, num_inference_steps=num_steps, generator=generator, ).images[0] return image @spaces.GPU(duration=170) def create_pipeline_logic(prompt_text, model_name): print(f"starting {model_name}") progress = gr.Progress(track_tqdm=True) num_steps = 30 guidance_scale = 7.5 # Example guidance scale, can be adjusted per model seed = 42 config = MODEL_CONFIGS[model_name] pipe_class = config["pipeline_class"] pipe = None if model_name == "AuraFlow": print("AuraFlow specal") b_pipe = AutoPipelineForText2Image.from_pretrained( config["repo_id"], variant="fp16", #cache_dir=config["cache_dir"], torch_dtype=torch.bfloat16 ).to("cuda") pipe = ProgressAuraFlowPipeline(b_pipe) else: pipe = AutoPipelineForText2Image.from_pretrained( config["repo_id"], variant="fp16", #cache_dir=config["cache_dir"], torch_dtype=torch.bfloat16 ).to("cuda") image = generate_image_with_progress( model_name,pipe, prompt_text, num_steps=num_steps, guidance_scale=guidance_scale, seed=seed, progress=progress ) return f"Seed: {seed}", image def main(): with gr.Blocks() as app: gr.Markdown("# Dynamic Multiple Model Image Generation") prompt_text = gr.Textbox(label="Enter prompt") for model_name, config in MODEL_CONFIGS.items(): with gr.Tab(model_name): button = gr.Button(f"Run {model_name}") output = gr.Textbox(label="Status") img = gr.Image(label=model_name, height=300) button.click(fn=create_pipeline_logic, inputs=[prompt_text, gr.Text(value= model_name,visible=False)], outputs=[output, img]) app.launch() if __name__ == "__main__": main()