app.py

import gradio as gr
import spaces
import torch
from diffusers import FluxKontextPipeline
from diffusers.utils import load_image
from PIL import Image
import os

# ----------------------------------------------
# Style → LoRA file mapping
# ----------------------------------------------
STYLE_LORA_MAP = {
    "3D_Chibi": "3D_Chibi_lora_weights.safetensors",
    "American_Cartoon": "American_Cartoon_lora_weights.safetensors",
    "Chinese_Ink": "Chinese_Ink_lora_weights.safetensors",
    "Clay_Toy": "Clay_Toy_lora_weights.safetensors",
    "Fabric": "Fabric_lora_weights.safetensors",
    "Ghibli": "Ghibli_lora_weights.safetensors",
    "Irasutoya": "Irasutoya_lora_weights.safetensors",
    "Jojo": "Jojo_lora_weights.safetensors",
    "Oil_Painting": "Oil_Painting_lora_weights.safetensors",
    "Pixel": "Pixel_lora_weights.safetensors",
    "Snoopy": "Snoopy_lora_weights.safetensors",
    "Poly": "Poly_lora_weights.safetensors",
    "LEGO": "LEGO_lora_weights.safetensors",
    "Origami": "Origami_lora_weights.safetensors",
    "Pop_Art": "Pop_Art_lora_weights.safetensors",
    "Van_Gogh": "Van_Gogh_lora_weights.safetensors",
    "Paper_Cutting": "Paper_Cutting_lora_weights.safetensors",
    "Line": "Line_lora_weights.safetensors",
    "Vector": "Vector_lora_weights.safetensors",
    "Picasso": "Picasso_lora_weights.safetensors",
    "Macaron": "Macaron_lora_weights.safetensors",
    "Rick_Morty": "Rick_Morty_lora_weights.safetensors"
}

# ----------------------------------------------
# Style descriptions (툴팁용)
# ----------------------------------------------
STYLE_DESCRIPTIONS = {
    "3D_Chibi": "귀여운 SD 캐릭터의 3D 느낌",
    "American_Cartoon": "고전적인 미국 카툰 스타일",
    "Chinese_Ink": "수묵화의 번짐과 농담 표현",
    "Clay_Toy": "찰흙·플라스틴 장난감質",
    "Fabric": "섬유·패브릭 질감",
    "Ghibli": "지브리풍 따뜻한 색감 & 연필선",
    "Irasutoya": "일러스토야 미니멀 평면 그림",
    "Jojo": "죠죠의 기묘한 모험 망가 터치",
    "Oil_Painting": "유화 붓터치와 질감",
    "Pixel": "16/32‑bit 레트로 픽셀아트",
    "Snoopy": "피너츠 스트립(스누피) 스타일",
    "Poly": "로우폴리 3D 기하학적 스타일",
    "LEGO": "레고 블록 조립 스타일",
    "Origami": "종이접기 질감·각도",
    "Pop_Art": "팝아트의 선명한 색감과 도트",
    "Van_Gogh": "반 고흐의 굵은 임파스토",
    "Paper_Cutting": "종이 오리기 실루엣",
    "Line": "깔끔한 라인 드로잉",
    "Vector": "벡터 그래픽·플랫 디자인",
    "Picasso": "피카소식 입체주의 큐비즘",
    "Macaron": "파스텔 마카롱톤 부드러움",
    "Rick_Morty": "릭 앤 모티 애니메이션 스타일"
}

# 전역 파이프라인
pipe = None


def get_dtype():
    """최적 dtype(bf16 지원 시 bf16, 아니면 fp16) 선택"""
    if torch.cuda.is_available():
        major, _ = torch.cuda.get_device_capability()
        if major >= 8:  # Ada/Hopper GPU는 bf16 본격 지원
            return torch.bfloat16
    return torch.float16


def load_pipeline():
    """FluxKontext 파이프라인을 (지연)로드"""
    global pipe
    if pipe is None:
        gr.Info("⬇️ FLUX.1‑Kontext 모델을 다운로드합니다…")

        pipe = FluxKontextPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-Kontext-dev",
            torch_dtype=get_dtype(),
            resume_download=True,
        )

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

        # VRAM 절약 모드
        if torch.cuda.is_available():
            pipe.enable_sequential_cpu_offload()
        pipe.vae.enable_tiling()

    return pipe


@spaces.GPU(duration=600)  # 초기 다운로드 시간 확보
def style_transfer(
    input_image,
    style_name,
    prompt_suffix,
    num_inference_steps,
    guidance_scale,
    seed,
):
    """선택한 스타일로 이미지 변환"""

    if input_image is None:
        gr.Warning("🖼️ 먼저 이미지를 업로드하세요!")
        return None

    try:
        pipe = load_pipeline()

        # 시드 고정 (0이면 난수)
        generator = None
        if seed and int(seed) != 0:
            generator = torch.Generator(device=pipe.device).manual_seed(int(seed))

        # 이미지 로드 & 리사이즈
        img = load_image(input_image) if isinstance(input_image, str) else input_image
        img = img.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS)

        # LoRA 로딩
        lora_file = STYLE_LORA_MAP[style_name]
        adapter_name = "style"
        pipe.load_lora_weights(
            "Owen777/Kontext-Style-Loras",
            weight_name=lora_file,
            adapter_name=adapter_name,
        )
        pipe.set_adapters([adapter_name], adapter_weights=[1.0])

        # 프롬프트 구성
        prompt = f"Turn this image into the {style_name.replace('_', ' ')} style."
        if prompt_suffix and prompt_suffix.strip():
            prompt += f" {prompt_suffix.strip()}"

        gr.Info("🎨 이미지를 생성 중…")

        result = pipe(
            image=img,
            prompt=prompt,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
            generator=generator,
            height=1024,
            width=1024,
        )

        # LoRA 언로드 및 캐시 정리
        pipe.unload_lora_weights()
        torch.cuda.empty_cache()

        return result.images[0]

    except Exception as e:
        gr.Error(f"🚨 오류: {e}")
        torch.cuda.empty_cache()
        return None


def update_description(style):
    return STYLE_DESCRIPTIONS.get(style, "")


with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 FLUX.1 Kontext Style Transfer
        FLUX.1‑Kontext‑dev 모델과 22개의 고품질 LoRA로 이미지를 다양한 예술 스타일로 변환하세요.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Upload Image", type="pil", height=400)

            style_dropdown = gr.Dropdown(
                choices=list(STYLE_LORA_MAP.keys()),
                value="Ghibli",
                label="Select Style",
                info="Choose from 22 different artistic styles",
            )

            style_info = gr.Textbox(
                label="Style Description",
                value=STYLE_DESCRIPTIONS["Ghibli"],
                interactive=False,
                lines=2,
            )

            prompt_suffix = gr.Textbox(
                label="Additional Instructions (Optional)",
                placeholder="예: 'make it more colorful', 'add dramatic lighting' …",
                lines=2,
            )

            with gr.Accordion("Advanced Settings", open=False):
                num_steps = gr.Slider(
                    10,
                    50,
                    value=24,
                    step=1,
                    label="Inference Steps",
                    info="더 높을수록 품질↑ 속도↓",
                )
                guidance = gr.Slider(
                    1.0,
                    7.5,
                    value=2.5,
                    step=0.1,
                    label="Guidance Scale",
                    info="프롬프트 준수 정도",
                )
                seed = gr.Number(label="Seed (0 = Random)", value=0)

            generate_btn = gr.Button(
                "🎨 Transform Image", variant="primary", size="lg"
            )

        with gr.Column(scale=1):
            output_image = gr.Image(label="Styled Result", type="pil", height=400)

            gr.Markdown(
                """### 💡 Tips:
                - 모든 이미지는 1024×1024로 리사이즈됩니다.
                - 첫 실행 시 7 GB 모델 다운로드가 필요합니다.
                - 스타일 변환은 약 30‑60 초 소요됩니다.
                - 다른 스타일도 시험해 보세요!"""
            )

    # 이벤트 바인딩
    style_dropdown.change(update_description, [style_dropdown], [style_info])

    generate_btn.click(
        style_transfer,
        inputs=[
            input_image,
            style_dropdown,
            prompt_suffix,
            num_steps,
            guidance,
            seed,
        ],
        outputs=[output_image],
    )

    gr.Markdown(
        """
        ---
        Created with ❤️ by [Black‑Forest Labs](https://huggingface.co/black-forest-labs) & 
        [Owen777/Kontext‑Style‑Loras](https://huggingface.co/Owen777/Kontext-Style-Loras)
        """
    )

if __name__ == "__main__":
    demo.launch()