python - How to show dedicated progress bar in each tab in a Gradio app? - Stack Overflow

admin管理员组

文章数量:1134577

I am developing an image generation Gradio app that uses multiple models like SD3.5, Flux, and others to generate images from a given prompt.

The app has 7 tabs, each corresponding to a specific model. Each tab displays an image generated by its respective model.

My problem is that I am unable to show a progress bar for each tab individually. Currently, the progress bar is displayed across all tabs simultaneously. However, I need a 'tab-specific progress bar.'

Below is my codebase and screenshots of the app that mocks the image generation process in the attachment. How can I implement this feature?

import random
from time import sleep

import gradio as gr
import threading
import requests
from PIL import Image
from io import BytesIO

# Constants
MAX_IMAGE_SIZE = 1024

# Model configurations
MODEL_CONFIGS = {
    "Stable Diffusion 3.5": {
        "repo_id": "stabilityai/stable-diffusion-3.5-large",
        "pipeline_class": "StableDiffusion3Pipeline"
    },
    "FLUX": {
        "repo_id": "black-forest-labs/FLUX.1-dev",
        "pipeline_class": "FluxPipeline"
    },
    "PixArt": {
        "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
        "pipeline_class": "PixArtSigmaPipeline"
    },
    "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"
    }
}

# Dictionary to store model pipelines
pipes = {}
model_locks = {model_name: threading.Lock() for model_name in MODEL_CONFIGS.keys()}


def fetch_image_from_url(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(BytesIO(response.content))
    except Exception as e:
        print(f"Error fetching image from URL {url}: {e}")
        return None


def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height,
                 guidance_scale, num_inference_steps):
    # Initialize a list to store all outputs
    all_outputs = [None] * (len(MODEL_CONFIGS) * 2)  # Pre-fill with None for each model's image and seed

    for idx, model_name in enumerate(MODEL_CONFIGS.keys()):
        try:
            progress_dict[model_name](0, desc=f"Starting generation for {model_name}...")
            print(f"IMAGE GENERATING {model_name}")
            generated_seed = seed if not randomize_seed else random.randint(0, 100000)

            # Fetch an image from a URL
            url = f".png?text=Hello+{model_name}+ +{generated_seed}"  # Replace with actual URL as needed
            image = fetch_image_from_url(url)

            progress_dict[model_name](0.9, desc=f"downloaded {model_name}...")
            # Update the outputs array with the result and seed, leaving remaining slots as None
            all_outputs[idx * 2] = image  # Image slot
            all_outputs[idx * 2 + 1] = generated_seed  # Seed slot

            # Add intermediate results to progress * (len(all_outputs) - len(all_outputs))
            yield all_outputs + [None]
            progress_dict[model_name](1, desc=f"generated {model_name}...")
            sleep(1)  # Simulate processing time

        except Exception as e:
            print(f"Error generating with {model_name}: {str(e)}")
            # Leave the slots for this model as None
            all_outputs[idx * 2] = None
            all_outputs[idx * 2 + 1] = None

    # Return the final completed array
    return all_outputs


# Gradio Interface
css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Multi-Model Image Generation")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Generate", scale=0, variant="primary")

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=100,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        memory_indicator = gr.Markdown("Current memory usage: 0 GB")

        with gr.Row():
            with gr.Column(scale=2):
                with gr.Tabs() as tabs:
                    results = {}
                    seeds = {}
                    progress_dict: dict[str, gr.Progress] = {}

                    for model_name in MODEL_CONFIGS.keys():
                        with gr.Tab(model_name):
                            results[model_name] = gr.Image(label=f"{model_name} Result")
                            seeds[model_name] = gr.Number(label="Seed used", visible=True)
                            progress_dict[model_name] = gr.Progress()

    # Prepare the input and output components
    input_components = [
        prompt, seed, randomize_seed,
    ]

    output_components = []
    for model_name in MODEL_CONFIGS.keys():
        output_components.extend([results[model_name], seeds[model_name]])

    run_button.click(
        fn=generate_all,
        inputs=input_components,
        outputs=output_components,
    )

if __name__ == "__main__":
    demo.launch()

I am developing an image generation Gradio app that uses multiple models like SD3.5, Flux, and others to generate images from a given prompt.

The app has 7 tabs, each corresponding to a specific model. Each tab displays an image generated by its respective model.

My problem is that I am unable to show a progress bar for each tab individually. Currently, the progress bar is displayed across all tabs simultaneously. However, I need a 'tab-specific progress bar.'

Below is my codebase and screenshots of the app that mocks the image generation process in the attachment. How can I implement this feature?

import random
from time import sleep

import gradio as gr
import threading
import requests
from PIL import Image
from io import BytesIO

# Constants
MAX_IMAGE_SIZE = 1024

# Model configurations
MODEL_CONFIGS = {
    "Stable Diffusion 3.5": {
        "repo_id": "stabilityai/stable-diffusion-3.5-large",
        "pipeline_class": "StableDiffusion3Pipeline"
    },
    "FLUX": {
        "repo_id": "black-forest-labs/FLUX.1-dev",
        "pipeline_class": "FluxPipeline"
    },
    "PixArt": {
        "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
        "pipeline_class": "PixArtSigmaPipeline"
    },
    "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"
    }
}

# Dictionary to store model pipelines
pipes = {}
model_locks = {model_name: threading.Lock() for model_name in MODEL_CONFIGS.keys()}


def fetch_image_from_url(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(BytesIO(response.content))
    except Exception as e:
        print(f"Error fetching image from URL {url}: {e}")
        return None


def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height,
                 guidance_scale, num_inference_steps):
    # Initialize a list to store all outputs
    all_outputs = [None] * (len(MODEL_CONFIGS) * 2)  # Pre-fill with None for each model's image and seed

    for idx, model_name in enumerate(MODEL_CONFIGS.keys()):
        try:
            progress_dict[model_name](0, desc=f"Starting generation for {model_name}...")
            print(f"IMAGE GENERATING {model_name}")
            generated_seed = seed if not randomize_seed else random.randint(0, 100000)

            # Fetch an image from a URL
            url = f"https://placehold.co/600x400/000000/FFFFFF.png?text=Hello+{model_name}+ +{generated_seed}"  # Replace with actual URL as needed
            image = fetch_image_from_url(url)

            progress_dict[model_name](0.9, desc=f"downloaded {model_name}...")
            # Update the outputs array with the result and seed, leaving remaining slots as None
            all_outputs[idx * 2] = image  # Image slot
            all_outputs[idx * 2 + 1] = generated_seed  # Seed slot

            # Add intermediate results to progress * (len(all_outputs) - len(all_outputs))
            yield all_outputs + [None]
            progress_dict[model_name](1, desc=f"generated {model_name}...")
            sleep(1)  # Simulate processing time

        except Exception as e:
            print(f"Error generating with {model_name}: {str(e)}")
            # Leave the slots for this model as None
            all_outputs[idx * 2] = None
            all_outputs[idx * 2 + 1] = None

    # Return the final completed array
    return all_outputs


# Gradio Interface
css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Multi-Model Image Generation")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Generate", scale=0, variant="primary")

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=100,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        memory_indicator = gr.Markdown("Current memory usage: 0 GB")

        with gr.Row():
            with gr.Column(scale=2):
                with gr.Tabs() as tabs:
                    results = {}
                    seeds = {}
                    progress_dict: dict[str, gr.Progress] = {}

                    for model_name in MODEL_CONFIGS.keys():
                        with gr.Tab(model_name):
                            results[model_name] = gr.Image(label=f"{model_name} Result")
                            seeds[model_name] = gr.Number(label="Seed used", visible=True)
                            progress_dict[model_name] = gr.Progress()

    # Prepare the input and output components
    input_components = [
        prompt, seed, randomize_seed,
    ]

    output_components = []
    for model_name in MODEL_CONFIGS.keys():
        output_components.extend([results[model_name], seeds[model_name]])

    run_button.click(
        fn=generate_all,
        inputs=input_components,
        outputs=output_components,
    )

if __name__ == "__main__":
    demo.launch()

• python
• artificial-intelligence
• stable-diffusion
• gradio
• image-generation

Share Improve this question Follow edited Jan 2 at 5:58 RagAnt asked Jan 1 at 7:20 RagAntRagAnt 1,00444 gold badges1919 silver badges4040 bronze badges 3

• Is this a web app? Which bit of your code represents the progress bar? – halfer Commented Jan 1 at 11:35
• You can consider making the most minimal code and asking. – redoc Commented Jan 4 at 12:58
• @redoc this is the minimum code – RagAnt Commented Jan 5 at 6:19

Add a comment  | 

1 Answer 1

Sorted by: Reset to default 1 +100

You need to manually pass the progress bar and maintain tabs differently. Also need to use callback_on_step_end with manual ones is needed.

Try this code:

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()

本文标签： pythonHow to show dedicated progress bar in each tab in a Gradio appStack Overflow