app.py

import torch
from diffusers import StableDiffusion3Pipeline, StableDiffusion2Pipeline, StableDiffusionXLBasePipeline
import gradio as gr
import os
import random
import transformers
import numpy as np
from transformers import T5Tokenizer, T5ForConditionalGeneration
import spaces

HF_TOKEN = os.getenv("HF_TOKEN")

if torch.cuda.is_available():
    device = "cuda"
    print("Using GPU")
else:
    device = "cpu"
    print("Using CPU")


MAX_SEED = np.iinfo(np.int32).max

# Initialize the pipelines for each sd model
sd3_medium_pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
)
sd3_medium_pipe.to(device)

sd2_1_pipe = StableDiffusion2Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16
)
sd2_1_pipe.to(device)

sdxl_pipe = StableDiffusionXLBasePipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
)
sdxl_pipe.to(device)

# superprompt-v1
tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1")
model = T5ForConditionalGeneration.from_pretrained(
    "roborovski/superprompt-v1", device_map="auto", torch_dtype="auto"
)
model.to(device)

# toggle visibility the enhanced prompt output
def update_visibility(enhance_prompt):
    return gr.update(visible=enhance_prompt)


# Define the image generation function for the Arena tab
@spaces.GPU(duration=80)
def generate_arena_images(
    prompt,
    enhance_prompt,
    negative_prompt,
    num_inference_steps,
    height,
    width,
    guidance_scale,
    seed,
    num_images_per_prompt,
    model_choice_1,
    model_choice_2,
    progress=gr.Progress(track_tqdm=True),
):
    if seed == 0:
        seed = random.randint(1, 2**32 - 1)

    if enhance_prompt:
        transformers.set_seed(seed)

        input_text = f"Expand the following prompt to add more detail: {prompt}"
        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)

        outputs = model.generate(
            input_ids,
            max_new_tokens=512,
            repetition_penalty=1.2,
            do_sample=True,
            temperature=0.7,
            top_p=1,
            top_k=50,
        )
        prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)

    generator = torch.Generator().manual_seed(seed)

    # Generate images for both models
    images_1 = generate_single_image(
        prompt,
        negative_prompt,
        num_inference_steps,
        height,
        width,
        guidance_scale,
        seed,
        num_images_per_prompt,
        model_choice_1,
        generator,
    )
    images_2 = generate_single_image(
        prompt,
        negative_prompt,
        num_inference_steps,
        height,
        width,
        guidance_scale,
        seed,
        num_images_per_prompt,
        model_choice_2,
        generator,
    )

    return images_1, images_2, prompt


# Helper function to generate images for a single model
def generate_single_image(
    prompt,
    negative_prompt,
    num_inference_steps,
    height,
    width,
    guidance_scale,
    seed,
    num_images_per_prompt,
    model_choice,
    generator,
):
    # Select the correct pipeline based on the model choice
    if model_choice == "sd3 medium":
        pipe = sd3_medium_pipe
    elif model_choice == "sd2.1":
        pipe = sd2_1_pipe
    elif model_choice == "sdxl":
        pipe = sdxl_pipe
    else:
        raise ValueError(f"Invalid model choice: {model_choice}")

    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        height=height,
        width=width,
        guidance_scale=guidance_scale,
        generator=generator,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    return output

# Define the image generation function for the Individual tab
@spaces.GPU(duration=80)
def generate_individual_image(
    prompt,
    enhance_prompt,
    negative_prompt,
    num_inference_steps,
    height,
    width,
    guidance_scale,
    seed,
    num_images_per_prompt,
    model_choice,
    progress=gr.Progress(track_tqdm=True),
):
    if seed == 0:
        seed = random.randint(1, 2**32 - 1)

    if enhance_prompt:
        transformers.set_seed(seed)

        input_text = f"Expand the following prompt to add more detail: {prompt}"
        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)

        outputs = model.generate(
            input_ids,
            max_new_tokens=512,
            repetition_penalty=1.2,
            do_sample=True,
            temperature=0.7,
            top_p=1,
            top_k=50,
        )
        prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)

    generator = torch.Generator().manual_seed(seed)

    output = generate_single_image(
        prompt,
        negative_prompt,
        num_inference_steps,
        height,
        width,
        guidance_scale,
        seed,
        num_images_per_prompt,
        model_choice,
        generator,
    )

    return output, prompt


# Create the Gradio interface
examples = [
    ["A white car racing fast to the moon.", True],
    ["A woman in a red dress singing on top of a building.", True],
    ["An astronaut on mars in a futuristic cyborg suit.", True],
]

css = """
.gradio-container{max-width: 1000px !important}
h1{text-align:center}
"""
with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column():
            gr.HTML(
                """
            <h1 style='text-align: center'>
            Stable Diffusion Arena
            </h1>
            """
            )
            gr.HTML(
                """
               Made by <a href='https://linktr.ee/Nick088' target='_blank'>Nick088</a>
               <br> <a href="https://discord.gg/osai"> <img src="https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge" alt="Discord"> </a>
                """
            )
    with gr.Tabs():
        with gr.TabItem("Arena"):
            with gr.Group():
                with gr.Column():
                    prompt = gr.Textbox(
                        label="Prompt",
                        info="Describe the image you want",
                        placeholder="A cat...",
                    )
                    enhance_prompt = gr.Checkbox(
                        label="Prompt Enhancement with SuperPrompt-v1", value=True
                    )
                    model_choice_1 = gr.Dropdown(
                        label="Stable Diffusion Model 1",
                        choices=["sd3 medium", "sd2.1", "sdxl"],
                        value="sd3 medium",
                    )
                    model_choice_2 = gr.Dropdown(
                        label="Stable Diffusion Model 2",
                        choices=["sd3 medium", "sd2.1", "sdxl"],
                        value="sd2.1",
                    )
                    run_button = gr.Button("Run")
                result_1 = gr.Gallery(label="Generated Images (Model 1)", elem_id="gallery_1")
                result_2 = gr.Gallery(label="Generated Images (Model 2)", elem_id="gallery_2")
                better_prompt = gr.Textbox(
                    label="Enhanced Prompt",
                    info="The output of your enhanced prompt used for the Image Generation",
                    visible=True,
                )
                enhance_prompt.change(
                    fn=update_visibility, inputs=enhance_prompt, outputs=better_prompt
                )
            with gr.Accordion("Advanced options", open=False):
                with gr.Row():
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        info="Describe what you don't want in the image",
                        value="deformed, distorted, disfigured, poorly drawn, bad anatomy, incorrect anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
                        placeholder="Ugly, bad anatomy...",
                    )
                with gr.Row():
                    num_inference_steps = gr.Slider(
                        label="Number of Inference Steps",
                        info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference",
                        minimum=1,
                        maximum=50,
                        value=25,
                        step=1,
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.",
                        minimum=0.0,
                        maximum=10.0,
                        value=7.5,
                        step=0.1,
                    )
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        info="Width of the Image",
                        minimum=256,
                        maximum=1344,
                        step=32,
                        value=1024,
                    )
                    height = gr.Slider(
                        label="Height",
                        info="Height of the Image",
                        minimum=256,
                        maximum=1344,
                        step=32,
                        value=1024,
                    )
                with gr.Row():
                    seed = gr.Slider(
                        value=42,
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        label="Seed",
                        info="A starting point to initiate the generation process, put 0 for a random one",
                    )
                    num_images_per_prompt = gr.Slider(
                        label="Images Per Prompt",
                        info="Number of Images to generate with the settings",
                        minimum=1,
                        maximum=4,
                        step=1,
                        value=2,
                    )

            gr.Examples(
                examples=examples,
                inputs=[prompt, enhance_prompt],
                outputs=[result_1, result_2, better_prompt],
                fn=generate_arena_images,
            )

            gr.on(
                triggers=[
                    prompt.submit,
                    run_button.click,
                ],
                fn=generate_arena_images,
                inputs=[
                    prompt,
                    enhance_prompt,
                    negative_prompt,
                    num_inference_steps,
                    width,
                    height,
                    guidance_scale,
                    seed,
                    num_images_per_prompt,
                    model_choice_1,
                    model_choice_2,
                ],
                outputs=[result_1, result_2, better_prompt],
            )

        with gr.TabItem("Individual"):
            with gr.Group():
                with gr.Column():
                    prompt = gr.Textbox(
                        label="Prompt",
                        info="Describe the image you want",
                        placeholder="A cat...",
                    )
                    enhance_prompt = gr.Checkbox(
                        label="Prompt Enhancement with SuperPrompt-v1", value=True
                    )
                    model_choice = gr.Dropdown(
                        label="Stable Diffusion Model",
                        choices=["sd3 medium", "sd2.1", "sdxl"],
                        value="sd3 medium",
                    )
                    run_button = gr.Button("Run")
                result = gr.Gallery(label="Generated AI Images", elem_id="gallery")
                better_prompt = gr.Textbox(
                    label="Enhanced Prompt",
                    info="The output of your enhanced prompt used for the Image Generation",
                    visible=True,
                )
                enhance_prompt.change(
                    fn=update_visibility, inputs=enhance_prompt, outputs=better_prompt
                )
            with gr.Accordion("Advanced options", open=False):
                with gr.Row():
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        info="Describe what you don't want in the image",
                        value="deformed, distorted, disfigured, poorly drawn, bad anatomy, incorrect anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
                        placeholder="Ugly, bad anatomy...",
                    )
                with gr.Row():
                    num_inference_steps = gr.Slider(
                        label="Number of Inference Steps",
                        info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference",
                        minimum=1,
                        maximum=50,
                        value=25,
                        step=1,
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.",
                        minimum=0.0,
                        maximum=10.0,
                        value=7.5,
                        step=0.1,
                    )
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        info="Width of the Image",
                        minimum=256,
                        maximum=1344,
                        step=32,
                        value=1024,
                    )
                    height = gr.Slider(
                        label="Height",
                        info="Height of the Image",
                        minimum=256,
                        maximum=1344,
                        step=32,
                        value=1024,
                    )
                with gr.Row():
                    seed = gr.Slider(
                        value=42,
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        label="Seed",
                        info="A starting point to initiate the generation process, put 0 for a random one",
                    )
                    num_images_per_prompt = gr.Slider(
                        label="Images Per Prompt",
                        info="Number of Images to generate with the settings",
                        minimum=1,
                        maximum=4,
                        step=1,
                        value=2,
                    )

            gr.Examples(
                examples=examples,
                inputs=[prompt, enhance_prompt],
                outputs=[result, better_prompt],
                fn=generate_individual_image,
            )

            gr.on(
                triggers=[
                    prompt.submit,
                    run_button.click,
                ],
                fn=generate_individual_image,
                inputs=[
                    prompt,
                    enhance_prompt,
                    negative_prompt,
                    num_inference_steps,
                    width,
                    height,
                    guidance_scale,
                    seed,
                    num_images_per_prompt,
                    model_choice,
                ],
                outputs=[result, better_prompt],
            )

demo.queue().launch(share=False)