app.py

import os
import random
import uuid
import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch
from diffusers import DiffusionPipeline

DESCRIPTION = """# Playground v2.5"""
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>CPU에서 실행 중 🥶 이 데모는 CPU에서 작동하지 않을 수 있습니다.</p>"

# 일반 설정
MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

# 디바이스 설정
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# CUDA가 사용 가능할 경우 글로벌로 파이프라인 로드
pipe = None
if torch.cuda.is_available():
    pipe = DiffusionPipeline.from_pretrained(
        "playgroundai/playground-v2.5-1024px-aesthetic",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16"
    ).to(device)

    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
    
    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

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

@spaces.GPU(enable_queue=True)
def generate(prompt: str, negative_prompt: str = "", use_negative_prompt: bool = False,
             seed: int = 0, width: int = 1024, height: int = 1024, guidance_scale: float = 7.0,
             randomize_seed: bool = False, use_resolution_binning: bool = True,
             progress=gr.Progress(track_tqdm=True)):
    global pipe
    if pipe is None:
        raise Exception("모델 파이프라인이 로드되지 않았습니다. 이 데모는 GPU가 필요합니다.")

    # pipe가 올바른 디바이스에 있는지 확인
    pipe.to(device)
    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator().manual_seed(seed)

    if not use_negative_prompt:
        negative_prompt = None
    
    images = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=25,
        generator=generator,
        num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
        use_resolution_binning=use_resolution_binning,
        output_type="pil",
    ).images

    image_paths = [save_image(img) for img in images]
    print(image_paths)
    return image_paths, seed