app.py

"""
스테이블 디퓨전 WebUI - 허깅페이스 스페이스용
Gradio 인터페이스 + REST API를 통한 이미지 생성 (txt2img + img2img 지원)
"""

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
from PIL import Image
import os
import gc
import io
import base64
from typing import Optional
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# 사용 가능한 모델 목록
MODELS = {
    "🎨 Mistoon Anime V3 (카툰풍 애니메이션)": "stablediffusionapi/mistoonanime-v30",
    "🌸 Anything V5 (애니메이션)": "stablediffusionapi/anything-v5",
    "💜 Counterfeit V3 (고품질 애니메이션)": "gsdf/Counterfeit-V3.0",
    "✨ DreamShaper V8 (다목적)": "Lykon/DreamShaper",
    "🎭 OpenJourney (Midjourney 스타일)": "prompthero/openjourney-v4",
    "🖼️ Stable Diffusion v1.5 (기본)": "runwayml/stable-diffusion-v1-5",
    "🌟 MeinaMix (애니메이션)": "Meina/MeinaMix_V11",
    "💫 ReV Animated (애니메이션)": "stablediffusionapi/rev-animated",
}

# 디바이스 설정
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32

print(f"🚀 디바이스: {DEVICE}, 데이터 타입: {DTYPE}")

# 현재 로드된 모델 정보
current_model_id = None
current_pipeline_type = None  # "txt2img" 또는 "img2img"
pipe = None

def clear_memory():
    """메모리 정리"""
    global pipe
    if pipe is not None:
        del pipe
        pipe = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_model(model_name: str, pipeline_type: str = "txt2img"):
    """
    모델 로드 함수
    
    Args:
        model_name: 모델 이름
        pipeline_type: "txt2img" 또는 "img2img"
    """
    global pipe, current_model_id, current_pipeline_type
    
    model_id = MODELS.get(model_name)
    if model_id is None:
        return None, f"❌ 알 수 없는 모델: {model_name}"
    
    # 이미 같은 모델과 파이프라인 타입이 로드되어 있으면 스킵
    if current_model_id == model_id and current_pipeline_type == pipeline_type and pipe is not None:
        return pipe, f"✅ {model_name} 이미 로드됨 ({pipeline_type})"
    
    # 기존 모델 정리
    clear_memory()
    
    print(f"📥 모델 로딩 중: {model_name} ({pipeline_type})...")
    
    try:
        # 파이프라인 타입에 따라 다른 클래스 사용
        if pipeline_type == "img2img":
            pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
                model_id,
                torch_dtype=DTYPE,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=False
            )
        else:
            pipe = StableDiffusionPipeline.from_pretrained(
                model_id,
                torch_dtype=DTYPE,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=False
            )
        
        # 빠른 스케줄러 사용
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        
        # 디바이스로 이동
        pipe = pipe.to(DEVICE)
        
        # 메모리 최적화
        pipe.enable_attention_slicing()
        if hasattr(pipe, 'enable_vae_slicing'):
            pipe.enable_vae_slicing()
        
        current_model_id = model_id
        current_pipeline_type = pipeline_type
        print(f"✅ 모델 로딩 완료: {model_name} ({pipeline_type})")
        return pipe, f"✅ {model_name} 로딩 완료!"
    
    except Exception as e:
        print(f"❌ 모델 로딩 실패: {e}")
        return None, f"❌ 모델 로딩 실패: {str(e)}"

# 기본 네거티브 프롬프트
DEFAULT_NEGATIVE = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"

def generate_txt2img(
    model_name: str,
    prompt: str,
    negative_prompt: str = "",
    num_inference_steps: int = 25,
    guidance_scale: float = 7.5,
    width: int = 512,
    height: int = 512,
    seed: int = -1,
    progress=gr.Progress()
):
    """텍스트 → 이미지 생성 함수"""
    
    global pipe
    
    if not prompt.strip():
        return None, "⚠️ 프롬프트를 입력해주세요!"
    
    # 모델 로드
    progress(0.1, desc="모델 로딩 중...")
    pipe, status = load_model(model_name, "txt2img")
    if pipe is None:
        return None, status
    
    # 네거티브 프롬프트 설정
    if negative_prompt.strip():
        full_negative = f"{negative_prompt}, {DEFAULT_NEGATIVE}"
    else:
        full_negative = DEFAULT_NEGATIVE
    
    # 시드 설정
    if seed == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    
    try:
        progress(0.3, desc="이미지 생성 중...")
        print(f"🎨 [txt2img] 이미지 생성 중... 프롬프트: {prompt[:50]}...")
        
        result = pipe(
            prompt=prompt,
            negative_prompt=full_negative,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            generator=generator
        )
        
        image = result.images[0]
        progress(1.0, desc="완료!")
        print("✅ [txt2img] 이미지 생성 완료!")
        
        return image, f"✅ 생성 완료! (시드: {seed})"
    
    except Exception as e:
        print(f"❌ [txt2img] 이미지 생성 실패: {e}")
        return None, f"❌ 이미지 생성 실패: {str(e)}"

def generate_img2img(
    model_name: str,
    input_image: Image.Image,
    prompt: str,
    negative_prompt: str = "",
    strength: float = 0.75,
    num_inference_steps: int = 25,
    guidance_scale: float = 7.5,
    seed: int = -1,
    progress=gr.Progress()
):
    """이미지 → 이미지 변환 함수"""
    
    global pipe
    
    if input_image is None:
        return None, "⚠️ 이미지를 업로드해주세요!"
    
    if not prompt.strip():
        return None, "⚠️ 프롬프트를 입력해주세요!"
    
    # 모델 로드 (img2img 파이프라인)
    progress(0.1, desc="모델 로딩 중...")
    pipe, status = load_model(model_name, "img2img")
    if pipe is None:
        return None, status
    
    # 네거티브 프롬프트 설정
    if negative_prompt.strip():
        full_negative = f"{negative_prompt}, {DEFAULT_NEGATIVE}"
    else:
        full_negative = DEFAULT_NEGATIVE
    
    # 시드 설정
    if seed == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    
    try:
        progress(0.3, desc="이미지 변환 중...")
        print(f"🖼️ [img2img] 이미지 변환 중... 프롬프트: {prompt[:50]}...")
        
        # 입력 이미지를 RGB로 변환하고 크기 조정
        input_image = input_image.convert("RGB")
        
        # 이미지 크기를 64의 배수로 조정 (SD 요구사항)
        w, h = input_image.size
        w = (w // 64) * 64
        h = (h // 64) * 64
        if w == 0:
            w = 512
        if h == 0:
            h = 512
        input_image = input_image.resize((w, h), Image.LANCZOS)
        
        result = pipe(
            prompt=prompt,
            image=input_image,
            negative_prompt=full_negative,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator
        )
        
        image = result.images[0]
        progress(1.0, desc="완료!")
        print("✅ [img2img] 이미지 변환 완료!")
        
        return image, f"✅ 변환 완료! (시드: {seed}, 강도: {strength})"
    
    except Exception as e:
        print(f"❌ [img2img] 이미지 변환 실패: {e}")
        return None, f"❌ 이미지 변환 실패: {str(e)}"

# Gradio 인터페이스 생성
def create_interface():
    """Gradio 웹 인터페이스 생성"""
    
    # 커스텀 CSS
    custom_css = """
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        max-width: 1200px !important;
    }
    
    .generate-btn {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        border: none !important;
        color: white !important;
        font-weight: bold !important;
        font-size: 1.2em !important;
        padding: 15px 30px !important;
        border-radius: 10px !important;
        transition: all 0.3s ease !important;
        width: 100% !important;
    }
    
    .generate-btn:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.5) !important;
    }
    
    .title {
        text-align: center;
        background: linear-gradient(135deg, #ff6b9d 0%, #c44569 50%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-size: 2.8em;
        font-weight: bold;
        margin-bottom: 5px;
    }
    
    .subtitle {
        text-align: center;
        color: #888;
        font-size: 1.1em;
        margin-bottom: 25px;
    }
    
    .model-dropdown {
        border: 2px solid #764ba2 !important;
        border-radius: 8px !important;
    }
    
    .output-image {
        border-radius: 12px !important;
        box-shadow: 0 4px 15px rgba(0,0,0,0.1) !important;
    }
    
    .status-box {
        background: linear-gradient(135deg, #f5f7fa 0%, #e4e8eb 100%);
        border-radius: 8px;
        padding: 10px;
        text-align: center;
    }
    """
    
    with gr.Blocks(css=custom_css, title="Stable Diffusion WebUI - Anime") as demo:
        # 헤더
        gr.HTML("""
            <div class="title">🌸 Anime Diffusion WebUI</div>
            <div class="subtitle">애니메이션 스타일 이미지 생성기 | Text-to-Image & Image-to-Image</div>
        """)
        
        # 탭으로 txt2img / img2img 분리
        with gr.Tabs():
            # ============================================
            # 탭 1: 텍스트 → 이미지 (txt2img)
            # ============================================
            with gr.TabItem("🎨 텍스트 → 이미지"):
                with gr.Row():
                    # 왼쪽: 입력 패널
                    with gr.Column(scale=1):
                        txt2img_model = gr.Dropdown(
                            label="🤖 모델 선택",
                            choices=list(MODELS.keys()),
                            value="🎨 Mistoon Anime V3 (카툰풍 애니메이션)",
                            elem_classes=["model-dropdown"]
                        )
                        
                        txt2img_prompt = gr.Textbox(
                            label="📝 프롬프트",
                            placeholder="1girl, anime style, beautiful, masterpiece, best quality",
                            lines=3
                        )
                        
                        txt2img_negative = gr.Textbox(
                            label="🚫 네거티브 프롬프트",
                            placeholder="추가할 네거티브 프롬프트 (기본값이 자동 적용됩니다)",
                            lines=2
                        )
                        
                        with gr.Row():
                            txt2img_width = gr.Slider(label="📐 너비", minimum=256, maximum=768, value=512, step=64)
                            txt2img_height = gr.Slider(label="📐 높이", minimum=256, maximum=768, value=768, step=64)
                        
                        with gr.Row():
                            txt2img_steps = gr.Slider(label="🔄 스텝 수", minimum=10, maximum=50, value=25, step=1)
                            txt2img_guidance = gr.Slider(label="🎯 CFG 스케일", minimum=1.0, maximum=15.0, value=7.0, step=0.5)
                        
                        txt2img_seed = gr.Number(label="🎲 시드 (-1 = 랜덤)", value=-1, precision=0)
                        
                        txt2img_btn = gr.Button("🚀 이미지 생성", elem_classes=["generate-btn"])
                        txt2img_status = gr.Textbox(label="📊 상태", value="프롬프트를 입력해주세요", interactive=False)
                    
                    # 오른쪽: 출력 패널
                    with gr.Column(scale=1):
                        txt2img_output = gr.Image(label="🖼️ 생성된 이미지", type="pil", elem_classes=["output-image"])
                
                # 예제
                gr.Examples(
                    examples=[
                        ["🎨 Mistoon Anime V3 (카툰풍 애니메이션)", "1girl, solo, colorful, vibrant colors, cartoon style, school uniform, masterpiece", ""],
                        ["🌸 Anything V5 (애니메이션)", "1girl, solo, long blue hair, cherry blossoms, detailed, masterpiece, best quality", ""],
                        ["💜 Counterfeit V3 (고품질 애니메이션)", "1girl, kimono, japanese garden, autumn leaves, ultra detailed, 8k", ""],
                    ],
                    inputs=[txt2img_model, txt2img_prompt, txt2img_negative],
                    label="💡 예제 프롬프트"
                )
                
                # 이벤트 핸들러
                txt2img_btn.click(
                    fn=generate_txt2img,
                    inputs=[txt2img_model, txt2img_prompt, txt2img_negative, txt2img_steps, txt2img_guidance, txt2img_width, txt2img_height, txt2img_seed],
                    outputs=[txt2img_output, txt2img_status]
                )
            
            # ============================================
            # 탭 2: 이미지 → 이미지 (img2img)
            # ============================================
            with gr.TabItem("🖼️ 이미지 → 이미지"):
                with gr.Row():
                    # 왼쪽: 입력 패널
                    with gr.Column(scale=1):
                        img2img_model = gr.Dropdown(
                            label="🤖 모델 선택",
                            choices=list(MODELS.keys()),
                            value="🎨 Mistoon Anime V3 (카툰풍 애니메이션)",
                            elem_classes=["model-dropdown"]
                        )
                        
                        img2img_input = gr.Image(
                            label="📤 입력 이미지 (실사, 스케치 등)",
                            type="pil",
                            height=200
                        )
                        
                        img2img_prompt = gr.Textbox(
                            label="📝 프롬프트 (변환할 스타일)",
                            placeholder="anime style, colorful, masterpiece, best quality",
                            lines=3
                        )
                        
                        img2img_negative = gr.Textbox(
                            label="🚫 네거티브 프롬프트",
                            placeholder="추가할 네거티브 프롬프트",
                            lines=2
                        )
                        
                        img2img_strength = gr.Slider(
                            label="💪 변환 강도 (0.0=원본 유지, 1.0=완전 변환)",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.75,
                            step=0.05
                        )
                        
                        with gr.Row():
                            img2img_steps = gr.Slider(label="🔄 스텝 수", minimum=10, maximum=50, value=25, step=1)
                            img2img_guidance = gr.Slider(label="🎯 CFG 스케일", minimum=1.0, maximum=15.0, value=7.0, step=0.5)
                        
                        img2img_seed = gr.Number(label="🎲 시드 (-1 = 랜덤)", value=-1, precision=0)
                        
                        img2img_btn = gr.Button("🚀 이미지 변환", elem_classes=["generate-btn"])
                        img2img_status = gr.Textbox(label="📊 상태", value="이미지와 프롬프트를 입력해주세요", interactive=False)
                    
                    # 오른쪽: 출력 패널
                    with gr.Column(scale=1):
                        img2img_output = gr.Image(label="🖼️ 변환된 이미지", type="pil", elem_classes=["output-image"])
                
                # img2img 가이드
                with gr.Accordion("📖 img2img 사용 가이드", open=False):
                    gr.Markdown("""
                    ### 🖼️ Image-to-Image 변환 가이드
                    
                    **사용 방법**:
                    1. 변환하고 싶은 이미지를 업로드합니다 (실사 사진, 스케치 등)
                    2. 원하는 스타일을 프롬프트로 입력합니다
                    3. 변환 강도를 조절합니다
                    
                    **변환 강도 (Strength) 설명**:
                    - `0.3` - 원본 이미지를 많이 유지 (미세한 스타일 변화)
                    - `0.5` - 균형 잡힌 변환
                    - `0.75` - 프롬프트에 더 충실한 변환 (권장)
                    - `1.0` - 거의 새로 생성 (원본 무시)
                    
                    **추천 프롬프트**:
                    - 실사 → 애니메이션: `anime style, colorful, masterpiece`
                    - 스케치 → 완성본: `detailed illustration, colored, vibrant`
                    - 지브리 스타일: `studio ghibli style, soft colors, fantasy`
                    """)
                
                # 이벤트 핸들러
                img2img_btn.click(
                    fn=generate_img2img,
                    inputs=[img2img_model, img2img_input, img2img_prompt, img2img_negative, img2img_strength, img2img_steps, img2img_guidance, img2img_seed],
                    outputs=[img2img_output, img2img_status]
                )
        
        # 푸터
        gr.HTML("""
            <div style="text-align: center; margin-top: 30px; padding: 20px; color: #888; border-top: 1px solid #eee;">
                <p>🌸 Powered by Diffusers & Gradio | 🤗 Hugging Face Spaces</p>
                <p style="font-size: 0.9em;">⚠️ CPU 모드에서는 이미지 생성에 2-5분 정도 소요됩니다.</p>
                <p style="font-size: 0.85em; color: #aaa;">첫 실행 시 모델 다운로드로 인해 시간이 더 걸릴 수 있습니다.</p>
            </div>
        """)
    
    return demo

# ================================
# REST API 엔드포인트 정의
# ================================

# API 요청/응답 모델 정의
class GenerateRequest(BaseModel):
    """텍스트 → 이미지 생성 요청"""
    prompt: str = Field(..., description="이미지 생성 프롬프트")
    model_name: str = Field(default="🎨 Mistoon Anime V3 (카툰풍 애니메이션)", description="사용할 모델 이름")
    negative_prompt: str = Field(default="", description="네거티브 프롬프트")
    num_inference_steps: int = Field(default=25, ge=10, le=50, description="추론 스텝 수")
    guidance_scale: float = Field(default=7.5, ge=1.0, le=15.0, description="CFG 스케일")
    width: int = Field(default=512, ge=256, le=768, description="이미지 너비")
    height: int = Field(default=512, ge=256, le=768, description="이미지 높이")
    seed: int = Field(default=-1, description="시드 값 (-1이면 랜덤)")

class Img2ImgRequest(BaseModel):
    """이미지 → 이미지 변환 요청"""
    image_base64: str = Field(..., description="입력 이미지 (Base64 인코딩)")
    prompt: str = Field(..., description="변환 프롬프트")
    model_name: str = Field(default="🎨 Mistoon Anime V3 (카툰풍 애니메이션)", description="사용할 모델 이름")
    negative_prompt: str = Field(default="", description="네거티브 프롬프트")
    strength: float = Field(default=0.75, ge=0.1, le=1.0, description="변환 강도")
    num_inference_steps: int = Field(default=25, ge=10, le=50, description="추론 스텝 수")
    guidance_scale: float = Field(default=7.5, ge=1.0, le=15.0, description="CFG 스케일")
    seed: int = Field(default=-1, description="시드 값 (-1이면 랜덤)")

class GenerateResponse(BaseModel):
    """이미지 생성 응답"""
    success: bool
    message: str
    image_base64: Optional[str] = None
    seed: Optional[int] = None

class ModelsResponse(BaseModel):
    """모델 목록 응답"""
    models: list[str]

# FastAPI 앱 생성
api_app = FastAPI(
    title="Anime Diffusion API",
    description="애니메이션 스타일 이미지 생성 REST API (txt2img + img2img)",
    version="2.0.0"
)

# CORS 설정 추가 (외부 호출 허용)
api_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@api_app.get("/api/models", response_model=ModelsResponse)
async def get_models():
    """사용 가능한 모델 목록 조회"""
    return ModelsResponse(models=list(MODELS.keys()))

@api_app.post("/api/generate", response_model=GenerateResponse)
async def api_generate_txt2img(request: GenerateRequest):
    """
    텍스트 → 이미지 생성 API
    
    프롬프트를 전달하면 Base64로 인코딩된 이미지를 반환합니다.
    """
    global pipe
    
    if not request.prompt.strip():
        raise HTTPException(status_code=400, detail="프롬프트를 입력해주세요")
    
    if request.model_name not in MODELS:
        raise HTTPException(status_code=400, detail=f"알 수 없는 모델입니다. 사용 가능한 모델: {list(MODELS.keys())}")
    
    # 모델 로드
    pipe, status = load_model(request.model_name, "txt2img")
    if pipe is None:
        raise HTTPException(status_code=500, detail=status)
    
    # 네거티브 프롬프트 설정
    if request.negative_prompt.strip():
        full_negative = f"{request.negative_prompt}, {DEFAULT_NEGATIVE}"
    else:
        full_negative = DEFAULT_NEGATIVE
    
    # 시드 설정
    seed = request.seed
    if seed == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    
    try:
        print(f"🎨 [API txt2img] 이미지 생성 중... 프롬프트: {request.prompt[:50]}...")
        
        result = pipe(
            prompt=request.prompt,
            negative_prompt=full_negative,
            num_inference_steps=request.num_inference_steps,
            guidance_scale=request.guidance_scale,
            width=request.width,
            height=request.height,
            generator=generator
        )
        
        image = result.images[0]
        
        # 이미지를 Base64로 인코딩
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        
        print(f"✅ [API txt2img] 이미지 생성 완료! (시드: {seed})")
        
        return GenerateResponse(
            success=True,
            message="이미지 생성 완료",
            image_base64=image_base64,
            seed=seed
        )
    
    except Exception as e:
        print(f"❌ [API txt2img] 이미지 생성 실패: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@api_app.post("/api/img2img", response_model=GenerateResponse)
async def api_generate_img2img(request: Img2ImgRequest):
    """
    이미지 → 이미지 변환 API
    
    Base64 이미지와 프롬프트를 전달하면 변환된 이미지를 반환합니다.
    """
    global pipe
    
    if not request.prompt.strip():
        raise HTTPException(status_code=400, detail="프롬프트를 입력해주세요")
    
    if request.model_name not in MODELS:
        raise HTTPException(status_code=400, detail=f"알 수 없는 모델입니다. 사용 가능한 모델: {list(MODELS.keys())}")
    
    # Base64 이미지 디코딩
    try:
        image_data = base64.b64decode(request.image_base64)
        input_image = Image.open(io.BytesIO(image_data)).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"이미지 디코딩 실패: {str(e)}")
    
    # 모델 로드 (img2img)
    pipe, status = load_model(request.model_name, "img2img")
    if pipe is None:
        raise HTTPException(status_code=500, detail=status)
    
    # 네거티브 프롬프트 설정
    if request.negative_prompt.strip():
        full_negative = f"{request.negative_prompt}, {DEFAULT_NEGATIVE}"
    else:
        full_negative = DEFAULT_NEGATIVE
    
    # 시드 설정
    seed = request.seed
    if seed == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    
    try:
        print(f"🖼️ [API img2img] 이미지 변환 중... 프롬프트: {request.prompt[:50]}...")
        
        # 이미지 크기를 64의 배수로 조정
        w, h = input_image.size
        w = (w // 64) * 64
        h = (h // 64) * 64
        if w == 0:
            w = 512
        if h == 0:
            h = 512
        input_image = input_image.resize((w, h), Image.LANCZOS)
        
        result = pipe(
            prompt=request.prompt,
            image=input_image,
            negative_prompt=full_negative,
            strength=request.strength,
            num_inference_steps=request.num_inference_steps,
            guidance_scale=request.guidance_scale,
            generator=generator
        )
        
        image = result.images[0]
        
        # 이미지를 Base64로 인코딩
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        
        print(f"✅ [API img2img] 이미지 변환 완료! (시드: {seed})")
        
        return GenerateResponse(
            success=True,
            message="이미지 변환 완료",
            image_base64=image_base64,
            seed=seed
        )
    
    except Exception as e:
        print(f"❌ [API img2img] 이미지 변환 실패: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@api_app.get("/api/health")
async def health_check():
    """서버 상태 확인"""
    return {
        "status": "healthy",
        "device": DEVICE,
        "model_loaded": current_model_id is not None,
        "pipeline_type": current_pipeline_type
    }

# ================================
# 메인 실행
# ================================
if __name__ == "__main__":
    print("🌸 Anime Diffusion WebUI + API 시작...")
    print("  - txt2img: 텍스트 → 이미지 생성")
    print("  - img2img: 이미지 → 이미지 변환")
    
    # Gradio 앱 생성
    demo = create_interface()
    
    # FastAPI에 Gradio 마운트
    app = gr.mount_gradio_app(api_app, demo, path="/")
    
    # uvicorn으로 통합 서버 실행
    import uvicorn
    print("📡 API 문서: http://localhost:7860/docs")
    print("🌐 웹 UI: http://localhost:7860/")
    uvicorn.run(app, host="0.0.0.0", port=7860)