# 🎨 Connect AI LAB — 무료 오픈소스 AI 이미지 생성 + 편집 (ZeroGPU)
#   🎨 생성 = Krea-2-Turbo (12B) · ✏️ 편집 = FLUX.1-Kontext(최고급, 토큰 있으면) / InstructPix2Pix(폴백)
#   유튜브(@CONNECT-AI-LAB) 홍보 · 디자인은 기본 테마 + 상단 배너.
import random

import gradio as gr
import numpy as np
import spaces
import torch

GEN_ID = "krea/Krea-2-Turbo"
EDIT_ID = "black-forest-labs/FLUX.2-klein-4B"   # 최신·ungated·Apache 2.0·생성+편집 통합(4B)
IP2P_ID = "timbrooks/instruct-pix2pix"          # 폴백(ungated)
YOUTUBE = "https://www.youtube.com/@CONNECT-AI-LAB"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = np.iinfo(np.int32).max

# ── 🎨 생성 (Krea-2-Turbo) ──
try:
    from diffusers import Krea2Pipeline
    gen_pipe = Krea2Pipeline.from_pretrained(GEN_ID, torch_dtype=torch.bfloat16)
except Exception:
    from diffusers import DiffusionPipeline
    gen_pipe = DiffusionPipeline.from_pretrained(GEN_ID, torch_dtype=torch.bfloat16, trust_remote_code=True)
gen_pipe = gen_pipe.to(DEVICE)

# ── ✏️ 편집: FLUX.2-klein-4B(최신·ungated) 우선, 실패 시 InstructPix2Pix 폴백 ──
EDIT_MODE = "ip2p"
try:
    from diffusers import Flux2KleinPipeline
    edit_pipe = Flux2KleinPipeline.from_pretrained(EDIT_ID, torch_dtype=torch.bfloat16)
    EDIT_MODE = "flux2"
except Exception as e:
    print(f"[edit] FLUX.2-klein 불가({e}); InstructPix2Pix로 폴백")
    from diffusers import StableDiffusionInstructPix2PixPipeline
    edit_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        IP2P_ID, torch_dtype=torch.float16, safety_checker=None)
edit_pipe = edit_pipe.to(DEVICE)


@spaces.GPU(duration=50)
def generate(prompt, seed, randomize_seed, steps, progress=gr.Progress(track_tqdm=True)):
    if not prompt or not prompt.strip():
        raise gr.Error("무엇을 그릴지 입력해주세요 🙂 (예: 눈 속의 귀여운 여우)")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    g = torch.Generator(device=DEVICE).manual_seed(int(seed))
    image = gen_pipe(prompt, num_inference_steps=int(steps), generator=g).images[0]
    return image, seed


@spaces.GPU(duration=55)
def edit(image, instruction, steps, strength, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("편집할 이미지를 먼저 올려주세요 🙂")
    if not instruction or not instruction.strip():
        raise gr.Error("어떻게 바꿀지 적어주세요 (예: 하늘을 보라색으로, 겨울로 바꿔줘)")
    img = image.convert("RGB")
    img.thumbnail((1024, 1024))
    if EDIT_MODE == "flux2":
        out = edit_pipe(image=img, prompt=instruction,
                        guidance_scale=float(strength), num_inference_steps=int(steps)).images[0]
    else:
        out = edit_pipe(instruction, image=img, num_inference_steps=int(steps),
                        image_guidance_scale=1.5, guidance_scale=7.5).images[0]
    return out


GEN_EXAMPLES = [
    "눈 속을 걷는 귀여운 아기 여우, 따뜻한 조명, 시네마틱, 고화질",
    "네온 사인 가득한 밤의 서울 거리, 비 오는 사이버펑크, 디테일",
    "미니멀한 로고: 초록빛 AI 두뇌, 심플, 벡터 스타일, 흰 배경",
    "우주에서 지구를 바라보는 우주비행사, 사실적, 4k, 영화 같은",
]
EDIT_EXAMPLES = ["하늘을 보라색 노을로 바꿔줘", "겨울 눈 내리는 장면으로", "수채화 그림 스타일로",
                 "네온 사이버펑크 느낌으로", "배경을 우주로 바꿔줘"]

CSS = """
.gradio-container {max-width: 960px !important; margin: auto !important;}
#hero {text-align:center; padding: 26px 16px; margin-bottom: 8px; border-radius: 18px;
  background: linear-gradient(135deg, #0c2a1a, #0a1830); border: 1px solid rgba(57,255,20,.25);}
#hero .brand {font-family: ui-monospace, monospace; letter-spacing: 6px; color: #58ff86; font-size: 12px;}
#hero h1 {color:#ffffff; font-size: 34px; font-weight: 800; margin: 8px 0 4px;}
#hero .sub {color:#bfeecf; font-size: 14px; margin: 2px 0 0;}
#hero .grn {color:#39ff14; font-weight:700;}
.yt-btn {display:inline-flex; align-items:center; gap:8px; margin-top:14px; padding:12px 24px; border-radius:12px;
  text-decoration:none; font-weight:800; font-size:14px; color:#ffffff !important;
  background:linear-gradient(135deg,#ff2d55,#ff0000); box-shadow:0 6px 22px rgba(255,0,0,.35);}
.yt-btn:hover {filter:brightness(1.08);}
#foot {text-align:center; color:#8a8f98; font-size:13px; padding:16px 8px;}
#foot a {color:#e11d48; text-decoration:none; font-weight:700;}
"""

_edit_note = ("✏️ 편집: FLUX.2-klein (최신·고급)" if EDIT_MODE == "flux2"
              else "✏️ 편집: 기본 모드 (FLUX.2 로드 대기)")

with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="green"), title="Connect AI LAB") as demo:
    gr.HTML(f"""
    <div id="hero">
      <div class="brand">◈ CONNECT · AI · LAB</div>
      <h1>🎨 AI 이미지 생성 &amp; 편집</h1>
      <div class="sub">텍스트로 <span class="grn">그리고</span>, 이미지를 <span class="grn">편집</span>까지 · 무료 · 오픈소스</div>
      <a class="yt-btn" href="{YOUTUBE}" target="_blank">▶ 유튜브 구독 · @CONNECT-AI-LAB</a>
    </div>
    """)

    with gr.Tabs():
        with gr.TabItem("🎨 생성"):
            with gr.Row():
                g_prompt = gr.Textbox(label="무엇을 그릴까요?",
                                      placeholder="예: 눈 속의 귀여운 아기 여우, 따뜻한 조명, 시네마틱",
                                      scale=5, lines=2)
                g_run = gr.Button("🎨 그리기", variant="primary", scale=1)
            g_result = gr.Image(label="결과", height=520)
            with gr.Accordion("⚙️ 세부 설정", open=False):
                with gr.Row():
                    g_steps = gr.Slider(4, 16, value=8, step=1, label="스텝")
                    g_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="시드")
                g_rand = gr.Checkbox(label="시드 랜덤(매번 다른 그림)", value=True)
            gr.Examples(examples=GEN_EXAMPLES, inputs=[g_prompt], label="✨ 예시")
            g_in = [g_prompt, g_seed, g_rand, g_steps]
            g_run.click(generate, inputs=g_in, outputs=[g_result, g_seed])
            g_prompt.submit(generate, inputs=g_in, outputs=[g_result, g_seed])

        with gr.TabItem("✏️ 편집"):
            gr.Markdown(f"##### {_edit_note}")
            with gr.Row():
                e_img = gr.Image(label="편집할 이미지 올리기", type="pil", height=340)
                e_out = gr.Image(label="편집 결과", height=340)
            with gr.Row():
                e_inst = gr.Textbox(label="어떻게 바꿀까요? (지시)",
                                    placeholder="예: 하늘을 보라색 노을로, 겨울로 바꿔줘, 수채화 스타일로",
                                    scale=5, lines=2)
                e_run = gr.Button("✏️ 편집", variant="primary", scale=1)
            with gr.Accordion("⚙️ 세부 설정", open=False):
                e_steps = gr.Slider(8, 40, value=18, step=1, label="스텝(낮을수록 빠름·쿼터 절약)")
                e_str = gr.Slider(1.0, 5.0, value=2.5, step=0.1, label="지시 반영 강도")
            gr.Examples(examples=EDIT_EXAMPLES, inputs=[e_inst], label="✨ 편집 예시")
            e_run.click(edit, inputs=[e_img, e_inst, e_steps, e_str], outputs=[e_out])

    gr.HTML(f"""
    <div id="foot">🌱 Powered by <b>Connect AI LAB</b> · 더 많은 AI 활용법은
      👉 <a href="{YOUTUBE}" target="_blank">유튜브 @CONNECT-AI-LAB</a> 구독!</div>
    """)

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