app.py

import os
import random
from typing import Callable, Dict, Optional, Tuple

import gradio as gr
import numpy as np
import PIL.Image
import spaces
import torch

from transformers import CLIPTextModel
from diffusers import AutoencoderKL, StableDiffusionXLPipeline, DDIMScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler

MODEL = "eienmojiki/Starry-XL-v5.2"
HF_TOKEN = os.getenv("HF_TOKEN")
MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
MAX_SEED = np.iinfo(np.int32).max

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

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

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        return seed

def seed_everything(seed: int) -> torch.Generator:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    generator = torch.Generator()
    generator.manual_seed(seed)
    return generator

def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
    scheduler_factory_map = {
        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
            scheduler_config, use_karras_sigmas=True
            ),
        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
            scheduler_config, use_karras_sigmas=True
            ),
        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
            scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
            ),
        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(scheduler_config),
        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
    }
    return scheduler_factory_map.get(name, lambda: None)()

def load_pipeline(model_name):
    pipe = StableDiffusionXLPipeline.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        safety_checker = None,
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
    )
    pipe.to(device)
    return pipe

@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = None,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 5.0,
    num_inference_steps: int = 24,
    sampler: str = "Euler a",
    clip_skip: int = 1,
    progress=gr.Progress(track_tqdm=True),
):
    
    generator = seed_everything(seed)
    pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)
    pipe.text_encoder = CLIPTextModel.from_pretrained(
        MODEL,
        subfolder = "text_encoder",
        num_hidden_layers = 12 - (clip_skip - 1),
        torch_dtype = torch.float16
    )

    try:

        img = pipe(
            prompt = prompt,
            negative_prompt = negative_prompt,
            width = width,
            height = height,
            guidance_scale = guidance_scale,
            num_inference_steps = num_inference_steps,
            generator = generator,
            output_type="pil",
        ).images[0]

        return img

    except Exception as e:
        print(f"An error occurred: {e}")

if torch.cuda.is_available():
    pipe = load_pipeline(MODEL)
    print("Loaded on Device!")
else:
    pipe = None