# ============================================================
# Open Model Playground
# Colab A100 Local + HF Space ZeroGPU FLUX.2 Fixed Edition
# Generated from Colab notebook cells.
# ============================================================

import os
import time
import random
import traceback
import gc
from typing import Tuple

import gradio as gr
from PIL import Image, ImageDraw

import torch

# `spaces`는 Hugging Face ZeroGPU에서는 필요하지만,
# Colab/local 환경에서는 설치되지 않았을 수 있습니다.
# import 실패 시 no-op fallback을 사용해 app.py import가 깨지지 않게 합니다.
try:
    import spaces
except Exception:
    class _SpacesFallback:
        def GPU(self, *args, **kwargs):
            def decorator(fn):
                return fn
            return decorator
    spaces = _SpacesFallback()

try:
    from huggingface_hub import login
    _hf_token = os.getenv("HF_TOKEN")
    if _hf_token:
        login(token=_hf_token)
        print("HF_TOKEN detected and login() called inside runtime.")
except Exception:
    traceback.print_exc()
APP_TITLE = "Open Model Playground"

TEXT_MODEL_ID = os.getenv("TEXT_MODEL_ID", "Qwen/Qwen3-4B-Instruct-2507")
IMAGE_MODEL_ID = os.getenv("IMAGE_MODEL_ID", "stabilityai/stable-diffusion-xl-base-1.0")
IMAGE_MODEL_DISPLAY = os.getenv("IMAGE_MODEL_DISPLAY", IMAGE_MODEL_ID)
SD_FALLBACK_MODEL_ID = os.getenv("SD_FALLBACK_MODEL_ID", "runwayml/stable-diffusion-v1-5")

DEVICE = "cuda" if torch.cuda.is_available() else (
    "mps" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else "cpu"
)

TEXT_MODEL = None
TEXT_TOKENIZER = None
IMAGE_PIPE = None

def _is_flux_model(model_id: str) -> bool:
    mid = (model_id or "").lower()
    return "flux" in mid or "black-forest-labs" in mid
def _now_ms() -> float:
    return time.time() * 1000


def _cleanup_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass


def _make_error_image(message: str, width: int = 512, height: int = 512) -> Image.Image:
    img = Image.new("RGB", (width, height), color=(255, 220, 220))
    draw = ImageDraw.Draw(img)
    draw.multiline_text(
        (18, 180),
        "이미지 생성 오류:\n" + str(message)[:500],
        fill=(180, 0, 0),
    )
    return img
def _get_text_model():
    global TEXT_MODEL, TEXT_TOKENIZER

    if TEXT_MODEL is not None and TEXT_TOKENIZER is not None:
        return TEXT_TOKENIZER, TEXT_MODEL

    from transformers import AutoModelForCausalLM, AutoTokenizer

    if DEVICE != "cuda":
        raise RuntimeError("실제 텍스트 모델 실행은 CUDA GPU를 권장합니다.")

    started = time.time()
    print(f"Loading text model: {TEXT_MODEL_ID}")

    TEXT_TOKENIZER = AutoTokenizer.from_pretrained(
        TEXT_MODEL_ID,
        trust_remote_code=True,
    )

    TEXT_MODEL = AutoModelForCausalLM.from_pretrained(
        TEXT_MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )

    TEXT_MODEL.eval()
    print(f"Loaded text model in {time.time() - started:.1f} sec")
    return TEXT_TOKENIZER, TEXT_MODEL
@spaces.GPU(duration=180)
def generate_text(
    system_message: str,
    user_prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
) -> Tuple[str, str]:
    start = time.time()

    try:
        system_message = system_message or "당신은 친절하고 정확한 AI 튜터입니다."
        user_prompt = user_prompt or ""

        if not user_prompt.strip():
            return "프롬프트를 입력해주세요.", "status=empty_prompt"

        tokenizer, model = _get_text_model()

        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_prompt},
        ]

        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
        else:
            text = f"{system_message}\n\nUser: {user_prompt}\nAssistant:"

        inputs = tokenizer(text, return_tensors="pt").to(model.device)
        do_sample = float(temperature) > 0

        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=int(max_new_tokens),
                temperature=float(temperature) if do_sample else None,
                top_p=float(top_p) if do_sample else None,
                do_sample=do_sample,
                pad_token_id=tokenizer.eos_token_id,
                use_cache=True,
            )

        generated = outputs[0][inputs["input_ids"].shape[-1]:]
        answer = tokenizer.decode(generated, skip_special_tokens=True)

        elapsed = time.time() - start
        meta = "\n".join([
            f"모델: {TEXT_MODEL_ID}",
            f"디바이스: {DEVICE}",
            f"소요 시간: {elapsed:.2f}초",
            f"max_new_tokens: {max_new_tokens}",
            f"temperature: {temperature}",
            f"top_p: {top_p}",
        ])
        return answer, meta

    except Exception as e:
        traceback.print_exc()
        elapsed = time.time() - start
        error_meta = "\n".join([
            f"오류 발생: {type(e).__name__}: {e}",
            f"소요 시간: {elapsed:.2f}초",
        ])
        return f"텍스트 생성 오류: {e}", error_meta
def _get_image_pipe():
    global IMAGE_PIPE

    if IMAGE_PIPE is not None:
        return IMAGE_PIPE

    start = time.time()

    if DEVICE == "cuda":
        from diffusers import DiffusionPipeline
        from diffusers import Flux2Pipeline

        print(f"Loading image model on CUDA: {IMAGE_MODEL_ID}")

        if _is_flux_model(IMAGE_MODEL_ID):
            # HF Space ZeroGPU FLUX.2 4-bit path.
            # 사용자가 공유한 기존 성공 코드 스타일에 맞춰 단순 DiffusionPipeline + to("cuda")를 사용합니다.
            IMAGE_PIPE = Flux2Pipeline.from_pretrained(
                IMAGE_MODEL_ID,
                torch_dtype=torch.bfloat16,
                trust_remote_code=True,
            )
            IMAGE_PIPE.to("cuda")
        else:
            # Colab A100 SDXL path.
            try:
                IMAGE_PIPE = DiffusionPipeline.from_pretrained(
                    IMAGE_MODEL_ID,
                    torch_dtype=torch.float16,
                    variant="fp16",
                    use_safetensors=True,
                )
            except Exception:
                IMAGE_PIPE = DiffusionPipeline.from_pretrained(
                    IMAGE_MODEL_ID,
                    torch_dtype=torch.float16,
                    use_safetensors=True,
                )
            IMAGE_PIPE.to("cuda")

        try:
            IMAGE_PIPE.set_progress_bar_config(disable=True)
        except Exception:
            pass

        print(f"Loaded image model in {time.time() - start:.1f} sec")
        return IMAGE_PIPE

    else:
        from diffusers import StableDiffusionPipeline

        print(f"Loading fallback image model on {DEVICE}: {SD_FALLBACK_MODEL_ID}")

        dtype = torch.float16 if DEVICE == "mps" else torch.float32

        IMAGE_PIPE = StableDiffusionPipeline.from_pretrained(
            SD_FALLBACK_MODEL_ID,
            torch_dtype=dtype,
        )

        IMAGE_PIPE.to(DEVICE)

        try:
            IMAGE_PIPE.set_progress_bar_config(disable=True)
        except Exception:
            pass

        print(f"Loaded fallback image model in {time.time() - start:.1f} sec")
        return IMAGE_PIPE
@spaces.GPU(duration=300)
def generate_image(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance_scale: float,
    seed: int,
    width: int,
    height: int,
) -> tuple:
    """이미지 생성 함수 — IMAGE_MODEL_ID에 따라 SDXL 또는 FLUX.2-dev 4-bit 사용"""
    start = time.time()

    try:
        prompt = prompt or "A friendly robot teaching open-source AI"
        negative_prompt = negative_prompt or ""

        pipe = _get_image_pipe()

        # CPU generator를 사용합니다.
        generator = torch.Generator(device="cpu").manual_seed(int(seed))

        if DEVICE == "cuda" and _is_flux_model(IMAGE_MODEL_ID):
            # FLUX.2-dev 4-bit, HF Space ZeroGPU
            result = pipe(
                prompt=prompt.strip(),
                num_inference_steps=int(steps),
                guidance_scale=float(guidance_scale),
                generator=generator,
                width=int(width),
                height=int(height),
            )
        elif DEVICE == "cuda":
            # SDXL on Colab A100
            result = pipe(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None,
                num_inference_steps=int(steps),
                guidance_scale=float(guidance_scale),
                generator=generator,
                width=int(width),
                height=int(height),
            )
        else:
            # SD v1.5 fallback on MPS/CPU
            result = pipe(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None,
                num_inference_steps=int(steps),
                guidance_scale=float(guidance_scale),
                generator=generator,
            )

        image = result.images[0]
        elapsed = time.time() - start
        meta = "\n".join([
            f"모델: {IMAGE_MODEL_DISPLAY}",
            f"모델 ID: {IMAGE_MODEL_ID}",
            f"디바이스: {DEVICE}",
            f"소요 시간: {elapsed:.2f}초",
            f"시드: {seed}",
            f"steps: {steps}",
            f"guidance_scale: {guidance_scale}",
            f"size: {width}x{height}",
        ])
        return image, meta

    except Exception as e:
        traceback.print_exc()
        try:
            img = _make_error_image(str(e))
            return img, f"오류 발생: {type(e).__name__}: {e}"
        except Exception:
            return None, f"오류 발생: {type(e).__name__}: {e}"
def build_demo():
    with gr.Blocks(title=APP_TITLE) as demo:
        gr.Markdown("# 🤗 Open Model Playground")
        gr.Markdown(
            "\n".join([
                f"**Text model:** `{TEXT_MODEL_ID}`  ",
                f"**Image model:** `{IMAGE_MODEL_ID}`  ",
                f"**Device:** `{DEVICE}`",
            ])
        )

        with gr.Tab("Text Generation"):
            gr.Markdown("## Text Generation with Qwen3 + Transformers")

            system_message = gr.Textbox(
                label="System message",
                value="당신은 친절하고 정확한 AI 튜터입니다.",
                lines=2,
            )
            user_prompt = gr.Textbox(
                label="User prompt",
                value="Hugging Face Hub가 AI 시대에 중요한 이유를 세 문장으로 설명해줘.",
                lines=4,
            )

            with gr.Row():
                max_new_tokens = gr.Slider(32, 1024, value=256, step=32, label="max_new_tokens")
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="temperature")
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")

            text_btn = gr.Button("Generate Text")
            text_output = gr.Textbox(label="Generated answer", lines=10)
            text_meta = gr.Textbox(label="Metadata", lines=7)

            text_btn.click(
                fn=generate_text,
                inputs=[
                    system_message,
                    user_prompt,
                    max_new_tokens,
                    temperature,
                    top_p,
                ],
                outputs=[text_output, text_meta],
                api_name="generate_text",
            )

        with gr.Tab("Image Generation"):
            gr.Markdown("## Image Generation with Diffusers")

            image_prompt = gr.Textbox(
                label="Prompt",
                value="A friendly robot teaching open-source AI in a modern Korean research lab, cinematic lighting, highly detailed",
                lines=3,
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                value="low quality, blurry, distorted, watermark, text artifacts",
                lines=2,
            )

            with gr.Row():
                steps = gr.Slider(4, 50, value=25, step=1, label="steps")
                guidance_scale = gr.Slider(0.0, 15.0, value=7.0, step=0.5, label="guidance_scale")
                seed = gr.Number(value=42, label="seed")

            with gr.Row():
                width = gr.Slider(512, 1024, value=1024, step=64, label="width")
                height = gr.Slider(512, 1024, value=1024, step=64, label="height")

            image_btn = gr.Button("Generate Image")
            image_output = gr.Image(label="Generated image")
            image_meta = gr.Textbox(label="Metadata", lines=8)

            image_btn.click(
                fn=generate_image,
                inputs=[
                    image_prompt,
                    negative_prompt,
                    steps,
                    guidance_scale,
                    seed,
                    width,
                    height,
                ],
                outputs=[image_output, image_meta],
                api_name="generate_image",
            )

        gr.Markdown(
            "\n".join([
                "### Runtime profile",
                "- Colab A100: Qwen3 + SDXL  ",
                "- HF Space ZeroGPU: Qwen3 + FLUX.2-dev-bnb-4bit",
            ])
        )

    return demo
def build_demo():
    with gr.Blocks(title=APP_TITLE) as demo:
        gr.Markdown("# 🤗 Open Model Playground")
        gr.Markdown(
            "\n".join([
                f"**Text model:** `{TEXT_MODEL_ID}`  ",
                f"**Image model:** `{IMAGE_MODEL_ID}`  ",
                f"**Device:** `{DEVICE}`",
            ])
        )

        with gr.Tab("Text Generation"):
            gr.Markdown("## Text Generation with Qwen3 + Transformers")

            system_message = gr.Textbox(
                label="System message",
                value="당신은 친절하고 정확한 AI 튜터입니다.",
                lines=2,
            )
            user_prompt = gr.Textbox(
                label="User prompt",
                value="Hugging Face Hub가 AI 시대에 중요한 이유를 세 문장으로 설명해줘.",
                lines=4,
            )

            with gr.Row():
                max_new_tokens = gr.Slider(32, 1024, value=256, step=32, label="max_new_tokens")
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="temperature")
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")

            text_btn = gr.Button("Generate Text")
            text_output = gr.Textbox(label="Generated answer", lines=10)
            text_meta = gr.Textbox(label="Metadata", lines=7)

            text_btn.click(
                fn=generate_text,
                inputs=[
                    system_message,
                    user_prompt,
                    max_new_tokens,
                    temperature,
                    top_p,
                ],
                outputs=[text_output, text_meta],
                api_name="generate_text",
            )

        with gr.Tab("Image Generation"):
            gr.Markdown("## Image Generation with Diffusers")

            image_prompt = gr.Textbox(
                label="Prompt",
                value="A friendly robot teaching open-source AI in a modern Korean research lab, cinematic lighting, highly detailed",
                lines=3,
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                value="low quality, blurry, distorted, watermark, text artifacts",
                lines=2,
            )

            with gr.Row():
                steps = gr.Slider(4, 50, value=25, step=1, label="steps")
                guidance_scale = gr.Slider(0.0, 15.0, value=7.0, step=0.5, label="guidance_scale")
                seed = gr.Number(value=42, label="seed")

            with gr.Row():
                width = gr.Slider(512, 1024, value=1024, step=64, label="width")
                height = gr.Slider(512, 1024, value=1024, step=64, label="height")

            image_btn = gr.Button("Generate Image")
            image_output = gr.Image(label="Generated image")
            image_meta = gr.Textbox(label="Metadata", lines=8)

            image_btn.click(
                fn=generate_image,
                inputs=[
                    image_prompt,
                    negative_prompt,
                    steps,
                    guidance_scale,
                    seed,
                    width,
                    height,
                ],
                outputs=[image_output, image_meta],
                api_name="generate_image",
            )

        gr.Markdown(
            "\n".join([
                "### Runtime profile",
                "- Colab A100: Qwen3 + SDXL  ",
                "- HF Space ZeroGPU: Qwen3 + FLUX.2-dev-bnb-4bit",
            ])
        )

    return demo
demo = build_demo()

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1)
    demo.launch()