app.py

import shutil

import gradio as gr

import torch
from diffusers import DiffusionPipeline

import wandb
from wandb.integration.diffusers import autolog


device = "cuda" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    torch.cuda.empty_cache()

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

if torch.cuda.is_available():
    pipeline.enable_xformers_memory_efficient_attention()


def generate_image(
    wandb_project,
    wandb_api_key,
    prompt,
    negative_prompt,
    height,
    width,
    num_inference_steps,
    guidance_scale,
    seed,
    guidance_rescale,
):
    if not (wandb_api_key is None or wandb_api_key == ""):
        wandb.login(key=wandb_api_key, relogin=True)
        generator = torch.Generator(device="cuda").manual_seed(seed)
        autolog(init={"project": wandb_project})
        run_url = wandb.run.get_url()
        image = pipeline(
            prompt,
            negative_prompt=negative_prompt,
            generator=generator,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
        ).images[0]
        wandb.finish()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        shutil.rmtree("wandb")
        return image, f"**WandB Run:** [{run_url}]({run_url})"
    else:
        return (
            None,
            "A WandB API key is required to run this app! You can get one from [https://wandb.ai/authorize](https://wandb.ai/authorize)",
        )


gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="WandB Project Name", value="Stable-Diffusion-XL"),
        gr.Textbox(
            label="WandB API Key (You can get one from https://wandb.ai/authorize)",
            type="password",
            value=None,
        ),
        gr.Textbox(
            label="Prompt",
            lines=3,
            value="a woman in orange, extremely detailed digital painting, in the style of Fenghua Zhong and Ruan Jia and jeremy lipking and Peter Mohrbacher, mystical colors, rim light, beautiful Lighting, 8k, stunning scene, raytracing, octane, trending on artstation",
        ),
        gr.Textbox(
            label="Negative Prompt",
            lines=3,
            value="blurry, plastic, grainy, duplicate, deformed , disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limb, disconnected limb, mutated hands and fingers, text, name, signature, watermark, worst quality, jpeg artifacts, boring composition, uninteresting",
        ),
        gr.Slider(minimum=512, maximum=1024, value=1024, step=128, label="Height"),
        gr.Slider(minimum=512, maximum=1024, value=1024, step=128, label="Width"),
        gr.Slider(
            minimum=10,
            maximum=100,
            value=50,
            step=10,
            label="Number of Inference Steps",
        ),
        gr.Slider(
            minimum=1.0,
            maximum=15.0,
            value=5.0,
            step=0.25,
            label="Guidance Scale (How Closely the model follows the Prompt)",
        ),
        gr.Slider(
            minimum=0, step=1, maximum=999999999999999999, randomize=True, label="Seed"
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.0,
            step=0.1,
            label="Guidance Rescale Factor",
        ),
    ],
    outputs=["image", "markdown"],
    title="Reproducible Stable Diffusion-XL with WandB",
    description="""
    This space enables us to generate images using Stable Diffusion XL and track your experiments using Weights & Biases. In order to learn more able engineering better prompts for Stable Diffusion models and how to make your experiments reproducible using [Weights & Biases](https://wandb.ai/site), check out [this report](https://wandb.ai/geekyrakshit/diffusers-prompt-engineering/reports/A-Guide-to-Prompt-Engineering-for-Diffusion-Models--Vmlldzo1NzY4NzQ3).
    This app is powered by [🤗 Diffusers](https://huggingface.co/docs/diffusers) and [Weights & Biases](https://wandb.ai/site). A **Weights & Biases API key is necessary** to run this app, which you can find at **[https://wandb.ai/authorize](https://wandb.ai/authorize)**.
    """,
    examples=[
        [
            "Stable-Diffusion-XL",
            None,
            "a woman in orange, extremely detailed digital painting, in the style of Fenghua Zhong and Ruan Jia and jeremy lipking and Peter Mohrbacher, mystical colors, rim light, beautiful Lighting, 8k, stunning scene, raytracing, octane, trending on artstation",
            "blurry, plastic, grainy, duplicate, deformed , disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limb, disconnected limb, mutated hands and fingers, text, name, signature, watermark, worst quality, jpeg artifacts, boring composition, uninteresting",
            1024,
            1024,
            50,
            5.0,
            256174057,
            0.0,
        ],
        [
            "Stable-Diffusion-XL",
            None,
            "hyper realistic photograph, Kodak disposable photography of young queen sit on her throne, blonde hair, red lips, white skin, 8k",
            None,
            1024,
            1024,
            50,
            5.0,
            189708808,
            0.0,
        ],
        [
            "Stable-Diffusion-XL",
            None,
            "hyper realistic waist up portrait of a teen boy with messy dark blond hair, light brown eyes, smiling, dimples, textured skin, realistic features, imperfect skin, 85mm lens Sony Lens, bokeh, sharp focus, professional photography, warm lighting, soft lights, realistic sharp eyes, sharp focus, set a busy street in the evening as background",
            None,
            1024,
            1024,
            50,
            5.0,
            267243010,
            0.0,
        ],
        [
            "Stable-Diffusion-XL",
            None,
            "A Cinematic Film Still Food Photography of a Cocktail, shot on fujifilm xt4, Professional Advertisement Photography, Dynamic Effects, Foaming Gas, Epic, Beautiful Details, Mysterious micro effects, Highly Detailed",
            None,
            1024,
            1024,
            50,
            5.0,
            1073711987,
            0.0,
        ],
        [
            "Stable-Diffusion-XL",
            None,
            "hyper realistic photograph, photography of a children dressed like a gangster, 50 mm, film grain, Kodak portra 800",
            None,
            1024,
            1024,
            50,
            5.0,
            355517910,
            0.0,
        ],
        [
            "Stable-Diffusion-XL",
            None,
            "a girl as personification of chocolate cupcake",
            "deformed, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limb, disconnected limb, mutated hands and fingers, blurry",
            1024,
            1024,
            50,
            5.0,
            595687707,
            0.0,
        ],
    ],
).launch(debug=True, max_threads=80)