Create app.py

app.py

@@ -0,0 +1,270 @@
+import spaces
+import torch
+import random
+import numpy as np
+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
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+class ProgressPipeline(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=3.5, seed=None,negative_prompt=None,  randomize_seed=None, width=1024, height=1024, num_inference_steps=40,  progress=gr.Progress(track_tqdm=True)):
+    generator = None
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    if seed is not None:
+        generator = torch.Generator("cuda").manual_seed(seed)
+    else:
+        generator = torch.Generator("cuda")
+
+    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
+    
+    # Define common arguments
+    common_args = {
+        "prompt": prompt,
+        "num_inference_steps": num_steps,
+        "negative_prompt": negative_prompt,
+        "width": width,
+        "height": height,
+        "generator": generator,
+    }
+    
+    if has_guidance_scale:
+        common_args["guidance_scale"] = guidance_scale
+    
+    if has_callback_on_step_end:
+        print("has callback_on_step_end and", "has guidance_scale" if has_guidance_scale else "NO guidance_scale")
+        common_args["callback_on_step_end"] = callback
+    else:
+        print("NO callback_on_step_end and", "has guidance_scale" if has_guidance_scale else "NO guidance_scale")
+        common_args["callback"] = callback
+        common_args["callback_steps"] = 1
+    
+    # Generate image
+    image = pipe(**common_args).images[0]
+
+    return seed, image
+
+@spaces.GPU(duration=170)
+def create_pipeline_logic(prompt_text, model_name, negative_prompt="",  seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=4.5, num_inference_steps=40,):
+    print(f"starting {model_name}")
+    progress = gr.Progress(track_tqdm=True)
+    config = MODEL_CONFIGS[model_name]
+    pipe_class = config["pipeline_class"]
+    pipe = None
+    b_pipe = AutoPipelineForText2Image.from_pretrained(
+        config["repo_id"],
+        #variant="fp16",
+        #cache_dir=config["cache_dir"],
+        torch_dtype=torch.bfloat16
+    ).to("cuda")
+    pipe_signature = signature(b_pipe)
+    # Check for the presence of "callback_on_step_end" in the signature
+    has_callback_on_step_end = "callback_on_step_end" in pipe_signature.parameters
+    if not has_callback_on_step_end:
+        pipe = ProgressPipeline(b_pipe)
+        print("ProgressPipeline specal")
+    else:
+        pipe = b_pipe
+        
+    gen_seed,image = generate_image_with_progress(
+        model_name,pipe, prompt_text, num_steps=num_inference_steps, guidance_scale=guidance_scale, seed=seed,negative_prompt = negative_prompt,  randomize_seed = randomize_seed, width = width, height = height, progress=progress
+    )
+    return f"Seed: {gen_seed}", image
+
+def main():
+    with gr.Blocks() as app:
+        gr.Markdown("# Dynamic Multiple Model Image Generation")
+
+        prompt_text = gr.Textbox(label="Enter prompt")
+
+        with gr.Accordion("Advanced Settings", open=False):
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+            )
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=100,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=7.5,
+                    step=0.1,
+                    value=4.5,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=40,
+                )
+
+        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), negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps], outputs=[output, img])
+
+    app.launch()
+
+
+if __name__ == "__main__":
+    main()
+