app.py

import shutil
from pathlib import Path

import gradio as gr
import requests
import spaces
import torch
from diffusers import (
    AutoencoderKL,
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionXLPipeline,
)
from loguru import logger
from PIL import Image
from tqdm import tqdm


def download(file: str, url: str):
    file_path = Path(file)
    if file_path.exists():
        return

    r = requests.get(url, stream=True)
    r.raise_for_status()

    temp_path = f"/tmp/{file_path.name}"
    with tqdm(
        desc=file, total=int(r.headers["content-length"]), unit="B", unit_scale=True
    ) as pbar, open(temp_path, "wb") as f:
        for chunk in r.iter_content(chunk_size=1024 * 1024):
            f.write(chunk)
            pbar.update(len(chunk))

    shutil.move(temp_path, file_path)


model_path = "pony-diffusion-v6-xl.safetensors"
download(
    model_path,
    "https://civitai.com/api/download/models/290640?type=Model&format=SafeTensor&size=pruned&fp=fp16",
)

vae_path = "pony-diffusion-v6-xl.vae.safetensors"
download(
    vae_path,
    "https://civitai.com/api/download/models/290640?type=VAE&format=SafeTensor",
)

vae = AutoencoderKL.from_single_file(
    vae_path,
    torch_dtype=torch.float16,
)

# pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
pipe = StableDiffusionXLPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16,
    vae=vae,
)
pipe = pipe.to("cuda")


@logger.catch(reraise=True)
@spaces.GPU
def generate(
    prompt: str,
    # init_image: Image.Image,
    strength: float,
    num_inference_steps: int,
    guidance_scale: float,
    progress=gr.Progress(track_tqdm=True),
):
    logger.info(f"Starting image generation: {dict(prompt=prompt, strength=strength)}")

    # # Downscale the image
    # init_image.thumbnail((1024, 1024))

    additional_args = {
        k: v
        for k, v in dict(
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        ).items()
        if v
    }

    images = pipe(
        prompt=prompt,
        # image=init_image,
        **additional_args,
    ).images

    return images[0]


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        # gr.Image(label="Init image", type="pil"),
        gr.Slider(label="Strength", minimum=0.0, maximum=1.0, value=0.0),
        gr.Slider(label="Number of inference steps", minimum=0, maximum=100, value=0),
        gr.Slider(label="Guidance scale", minimum=0.0, maximum=100.0, value=0.0),
    ],
    outputs=[gr.Image(label="Output")],
)

demo.launch()