app.py

"""
FLUX.2 [klein] 9B — Character Consistency Space
================================================
• Tab 1 │ Karakter Referansı   : Prompt + max 4 referans görsel → yeni görsel
• Tab 2 │ img2img + Referans   : Giriş görseli + prompt + max 3 referans → düzenlenmiş görsel

NOTLAR:
- Bu model, multi-reference desteğini NATIVE olarak sunar; ayrı IP-Adapter gerekmez.
- FLUX Non-Commercial License geçerlidir. Space'inize HF_TOKEN secret'ı eklemeyi unutmayın.
- Gereksinim: ~29 GB VRAM (A100 / H100 önerilir).
"""
import spaces
import os
import random
import gradio as gr
import numpy as np
import torch
from diffusers import Flux2KleinPipeline
from PIL import Image

# ──────────────────────────────────────────────
# Sabitler
# ──────────────────────────────────────────────

MODEL_ID     = "black-forest-labs/FLUX.2-klein-9B"
DTYPE        = torch.bfloat16
MAX_REFS_T2I = 4   # karakter referansı tabı için maks referans görseli
MAX_REFS_I2I = 3   # img2img tabı için maks referans görseli

# ──────────────────────────────────────────────
# Model yükleme (lazy, singleton)
# ──────────────────────────────────────────────

_pipe: Flux2KleinPipeline | None = None


def get_pipe() -> Flux2KleinPipeline:
    global _pipe
    if _pipe is None:
        print(f"[INFO] Model yükleniyor: {MODEL_ID}")
        token = os.environ.get("HF_TOKEN")          # HF Space secret
        _pipe = Flux2KleinPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=DTYPE,
            token=token,
        )
        #_pipe.enable_model_cpu_offload()             # VRAM'ı verimli kullan
        print("[INFO] Model hazır.")
    return _pipe


# ──────────────────────────────────────────────
# Yardımcı fonksiyonlar
# ──────────────────────────────────────────────

def to_pil(img) -> Image.Image | None:
    """Gradio'dan gelen numpy array veya PIL görselini PIL'e çevirir."""
    if img is None:
        return None
    if isinstance(img, np.ndarray):
        return Image.fromarray(img).convert("RGB")
    if isinstance(img, Image.Image):
        return img.convert("RGB")
    return None


def collect_refs(*images) -> list[Image.Image] | None:
    """None olmayan görselleri filtreler, PIL listesi döndürür."""
    refs = [to_pil(img) for img in images if img is not None]
    return refs if refs else None


def make_generator(seed: int) -> torch.Generator | None:
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    gen = torch.Generator()
    gen.manual_seed(seed)
    return gen


# ──────────────────────────────────────────────
# Tab 1 — Karakter Referansı (Multi-Ref T2I)
# ──────────────────────────────────────────────
@spaces.GPU
def run_t2i(
    prompt: str,
    ref1, ref2, ref3, ref4,
    width: int,
    height: int,
    steps: int,
    guidance: float,
    seed: int,
) -> Image.Image:
    if not prompt.strip():
        raise gr.Error("Lütfen bir prompt girin.")

    refs = collect_refs(ref1, ref2, ref3, ref4)
    gen  = make_generator(seed)
    pipe = get_pipe()

    result = pipe(
        prompt=prompt,
        image=refs,          # None → saf t2i | list → multi-ref yönlendirmeli üretim
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=guidance,
        generator=gen,
    ).images[0]

    return result


# ──────────────────────────────────────────────
# Tab 2 — img2img + Karakter Referansı
# ──────────────────────────────────────────────
@spaces.GPU
def run_i2i(
    prompt: str,
    input_img,
    ref1, ref2, ref3,
    width: int,
    height: int,
    steps: int,
    guidance: float,
    strength: float,
    seed: int,
) -> Image.Image:
    if not prompt.strip():
        raise gr.Error("Lütfen bir prompt girin.")
    if input_img is None:
        raise gr.Error("Lütfen düzenlemek istediğiniz bir giriş görseli yükleyin.")

    base = to_pil(input_img)
    refs = collect_refs(ref1, ref2, ref3)

    # Tüm görseller tek liste olarak verilir:
    # [giriş görseli] + [karakter referansları]
    # Model, prompt'a göre hangi unsurları transfer edeceğine karar verir.
    all_images = ([base] + refs) if refs else [base]

    gen  = make_generator(seed)
    pipe = get_pipe()

    # strength parametresi img2img gürültü oranını kontrol eder.
    # Flux2KleinPipeline bu parametreyi desteklemiyorsa sessizce atlanır.
    call_kwargs = dict(
        prompt=prompt,
        image=all_images,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=guidance,
        generator=gen,
    )
    try:
        result = pipe(**call_kwargs, strength=strength).images[0]
    except TypeError:
        # strength desteklenmiyorsa (pipeline versiyonuna göre değişebilir)
        result = pipe(**call_kwargs).images[0]

    return result


# ──────────────────────────────────────────────
# Gradio Arayüzü
# ──────────────────────────────────────────────

CSS = """
h1 { text-align: center; }
.ref-row { gap: 8px !important; }
footer { display: none !important; }
"""

# Paylaşılan çözünürlük seçenekleri
RESOLUTION_CHOICES = [512, 768, 1024, 1280]

with gr.Blocks(title="FLUX.2 [klein] Character Studio") as demo:

    gr.Markdown(
        "# 🎭 FLUX.2 \\[klein\\] 9B — Character Consistency Studio\n"
        "> Multi-reference karakter tutarlılığı · Native multi-image support · 4-step distilled"
    )

    # ── Tab 1: Karakter Referansı ──────────────────────────────────────────
    with gr.Tab("🧑‍🎨 Karakter Referansı (T2I)"):
        gr.Markdown(
            "Karakterinize ait **1–4 referans görsel** yükleyin. "
            "Model bu görsellerdeki kimliği (yüz, kıyafet, stil) yeni bir sahneye taşır."
        )
        with gr.Row():
            with gr.Column(scale=1):
                t2i_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="örn. The character in a snowy forest, cinematic lighting",
                    lines=3,
                )
                gr.Markdown("**Referans Görseller** (karakter kimliği için, maks. 4)")
                with gr.Row(elem_classes="ref-row"):
                    t2i_ref1 = gr.Image(label="Ref 1", type="pil", height=140)
                    t2i_ref2 = gr.Image(label="Ref 2", type="pil", height=140)
                    t2i_ref3 = gr.Image(label="Ref 3", type="pil", height=140)
                    t2i_ref4 = gr.Image(label="Ref 4", type="pil", height=140)

                with gr.Accordion("⚙️ Gelişmiş Ayarlar", open=False):
                    with gr.Row():
                        t2i_width  = gr.Dropdown(RESOLUTION_CHOICES, value=1024, label="Genişlik")
                        t2i_height = gr.Dropdown(RESOLUTION_CHOICES, value=1024, label="Yükseklik")
                    t2i_steps    = gr.Slider(1, 8,   value=4,   step=1,   label="Adım Sayısı (Steps)")
                    t2i_guidance = gr.Slider(0.0, 5.0, value=1.0, step=0.1, label="Guidance Scale")
                    t2i_seed     = gr.Number(value=-1, label="Seed (-1 = rastgele)", precision=0)

                t2i_btn = gr.Button("🎨 Oluştur", variant="primary", size="lg")

            with gr.Column(scale=1):
                t2i_out = gr.Image(label="Çıktı", type="pil", height=520)

        t2i_btn.click(
            fn=run_t2i,
            inputs=[
                t2i_prompt,
                t2i_ref1, t2i_ref2, t2i_ref3, t2i_ref4,
                t2i_width, t2i_height,
                t2i_steps, t2i_guidance, t2i_seed,
            ],
            outputs=t2i_out,
        )

    # ── Tab 2: img2img + Referans ──────────────────────────────────────────
    with gr.Tab("🖼️ img2img + Referans"):
        gr.Markdown(
            "Bir **giriş görseli** ve opsiyonel **karakter referansları** yükleyin. "
            "Model, giriş görselini prompt doğrultusunda düzenlerken karakteri referanslardan alır."
        )
        with gr.Row():
            with gr.Column(scale=1):
                i2i_prompt    = gr.Textbox(
                    label="Prompt",
                    placeholder="örn. The character wearing a red jacket, studio lighting",
                    lines=3,
                )
                i2i_input_img = gr.Image(label="📥 Giriş Görseli (düzenlenecek)", type="pil", height=220)
                gr.Markdown("**Karakter Referansları** (opsiyonel, maks. 3)")
                with gr.Row(elem_classes="ref-row"):
                    i2i_ref1 = gr.Image(label="Ref 1", type="pil", height=140)
                    i2i_ref2 = gr.Image(label="Ref 2", type="pil", height=140)
                    i2i_ref3 = gr.Image(label="Ref 3", type="pil", height=140)
 
                with gr.Accordion("⚙️ Gelişmiş Ayarlar", open=False):
                    i2i_strength = gr.Slider(
                        0.1, 1.0, value=0.75, step=0.05,
                        label="Strength (düşük = giriş görseliyle daha yakın)",
                    )
                    with gr.Row():
                        i2i_width  = gr.Dropdown(RESOLUTION_CHOICES, value=1024, label="Genişlik")
                        i2i_height = gr.Dropdown(RESOLUTION_CHOICES, value=1024, label="Yükseklik")
                    i2i_steps    = gr.Slider(1, 8,   value=4,   step=1,   label="Adım Sayısı (Steps)")
                    i2i_guidance = gr.Slider(0.0, 5.0, value=1.0, step=0.1, label="Guidance Scale")
                    i2i_seed     = gr.Number(value=-1, label="Seed (-1 = rastgele)", precision=0)

                i2i_btn = gr.Button("🖌️ Düzenle", variant="primary", size="lg")

            with gr.Column(scale=1):
                i2i_out = gr.Image(label="Çıktı", type="pil", height=520)

        i2i_btn.click(
            fn=run_i2i,
            inputs=[
                i2i_prompt,
                i2i_input_img,
                i2i_ref1, i2i_ref2, i2i_ref3,
                i2i_width, i2i_height,
                i2i_steps, i2i_guidance, i2i_strength, i2i_seed,
            ],
            outputs=i2i_out,
        )

    # ── Bilgi Notu ──────────────────────────────────────────────────────────
    gr.Markdown(
        "---\n"
        "**Model:** `black-forest-labs/FLUX.2-klein-9B` · "
        "**Lisans:** FLUX Non-Commercial · "
        "**Gereksinim:** ~29 GB VRAM (A100 / H100)"
    )

# ──────────────────────────────────────────────
# Başlatma
# ──────────────────────────────────────────────

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