"""
AI Paper Lab — ICML 2026 논문 체험 Space
pluto2060/ai-paper-lab
"""
import gradio as gr
import numpy as np
import plotly.graph_objects as go
import spaces
import torch

# ── Tab 1: Grok Grokking ──────────────────────────────────────────────────────

def grokking_simulation(n_train: int, weight_decay: float, n_epochs: int):
    """Simulate grokking in ridge regression — memorization then generalization."""
    np.random.seed(42)
    # Ground truth weight vector
    w_true = np.array([1.0, -1.5, 0.8])
    # Training data
    X_train = np.random.randn(n_train, 3)
    y_train = X_train @ w_true + np.random.randn(n_train) * 0.1
    # Test data (larger)
    X_test = np.random.randn(200, 3)
    y_test = X_test @ w_true + np.random.randn(200) * 0.1

    train_loss, test_loss, weight_norm = [], [], []
    epochs = np.arange(1, n_epochs + 1)

    # Gradient descent simulation
    w = np.zeros(3)
    lr = 0.01
    for ep in epochs:
        grad = -2 * X_train.T @ (y_train - X_train @ w) / n_train + 2 * weight_decay * w
        w = w - lr * grad
        tr = np.mean((y_train - X_train @ w) ** 2)
        te = np.mean((y_test - X_test @ w) ** 2)
        train_loss.append(tr)
        test_loss.append(te)
        weight_norm.append(np.linalg.norm(w))

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=epochs, y=train_loss, name="훈련 손실", line=dict(color="#2196F3")))
    fig.add_trace(go.Scatter(x=epochs, y=test_loss, name="테스트 손실", line=dict(color="#F44336")))
    fig.add_trace(go.Scatter(x=epochs, y=weight_norm, name="가중치 노름", line=dict(color="#4CAF50", dash="dot")))

    # Mark grokking transition: first epoch where test_loss < 0.5
    grok_ep = next((i + 1 for i, v in enumerate(test_loss) if v < 0.5), None)
    if grok_ep:
        fig.add_vline(x=grok_ep, line_dash="dash", line_color="orange",
                      annotation_text=f"그로킹 발생! (epoch {grok_ep})", annotation_position="top right")

    fig.update_layout(
        title="To Grok Grokking — 릿지 회귀의 지연 일반화 시뮬레이션",
        xaxis_title="Epoch", yaxis_title="값",
        template="plotly_dark", height=450
    )
    return fig


def build_grokking_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 To Grok Grokking: Provable Grokking in Ridge Regression
**ICML 2026 Honorable Mention** | Xu et al.

> AI 모델이 처음에는 훈련 데이터를 **암기**하다가 어느 순간 갑자기 **진짜 원리를 일반화**하는 현상을 수학적으로 증명.
> 릿지 회귀라는 단순 모델로 이 "그로킹" 현상이 왜 발생하는지 해석적으로 규명.

**체험**: 슬라이더를 조절해 그로킹이 언제 발생하는지 직접 확인하세요.
- `훈련 데이터 수` ↓ → 암기 기간 더 길어짐
- `정규화 강도` ↑ → 일반화 촉진 (그로킹 가속)
""")
        with gr.Row():
            n_train = gr.Slider(5, 50, value=15, step=5, label="훈련 데이터 수 (N)")
            wd = gr.Slider(0.001, 0.1, value=0.01, step=0.001, label="정규화 강도 (weight_decay)")
            n_ep = gr.Slider(100, 2000, value=500, step=100, label="학습 epoch 수")
        btn = gr.Button("시뮬레이션 실행", variant="primary")
        plot = gr.Plot()
        btn.click(grokking_simulation, inputs=[n_train, wd, n_ep], outputs=plot)


# ── Tab 2: Quantum vs Classical ───────────────────────────────────────────────

def quantum_classical_viz(n_max: int):
    ns = np.arange(1, n_max + 1)
    classical = ns ** 2          # O(n²) — typical classical linear system solver
    quantum_inspired = ns * np.log(ns) ** 2  # quantum-inspired classical (best known)
    quantum = np.log(ns) ** 2   # quantum algorithm (polylog)

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=ns, y=classical, name="고전 알고리즘 O(n²)", line=dict(color="#F44336")))
    fig.add_trace(go.Scatter(x=ns, y=quantum_inspired, name="양자영감 고전 O(n·log²n)", line=dict(color="#FF9800")))
    fig.add_trace(go.Scatter(x=ns, y=quantum, name="양자 알고리즘 O(log²n)", line=dict(color="#4CAF50")))

    fig.update_layout(
        title=f"양자 vs 고전 알고리즘 복잡도 비교 (N=1~{n_max})",
        xaxis_title="문제 크기 N", yaxis_title="계산 단계 수 (상대적)",
        template="plotly_dark", height=450,
        yaxis_type="log"
    )
    return fig


def build_quantum_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 An Exponential Separation Between Quantum and Classical Algorithms
**ICML 2026 Spotlight** | Grønlund & Larsen (Aarhus Univ.)

> 선형 방정식 풀기에서 양자 알고리즘이 고전 알고리즘보다 **지수함수적으로 빠름**을 수학적으로 증명.
> "지수 분리"란: 고전 컴퓨터가 백만 년 걸릴 계산을 양자 컴퓨터가 단 몇 초에 처리.

**체험**: N 슬라이더를 늘릴수록 세 알고리즘의 격차가 얼마나 벌어지는지 확인하세요.
""")
        n_max = gr.Slider(10, 500, value=100, step=10, label="문제 크기 N (최대)")
        btn = gr.Button("비교 차트 생성", variant="primary")
        plot = gr.Plot()
        btn.click(quantum_classical_viz, inputs=[n_max], outputs=plot)


# ── Tab 3: M+Adam vs Adam ─────────────────────────────────────────────────────

def madam_simulation(lr: float, beta1: float, noise_scale: float, n_steps: int):
    np.random.seed(0)
    # Simple 2D quadratic loss: f(x,y) = x² + 10y²
    def loss_fn(w): return w[0] ** 2 + 10 * w[1] ** 2
    def grad_fn(w): return np.array([2 * w[0], 20 * w[1]])

    def simulate(optimizer="adam"):
        w = np.array([3.0, 2.0])
        m, v = np.zeros(2), np.zeros(2)
        beta2 = 0.999
        eps = 1e-8
        losses = []
        for t in range(1, n_steps + 1):
            g = grad_fn(w) + np.random.randn(2) * noise_scale
            if optimizer == "adam":
                m = beta1 * m + (1 - beta1) * g
                v = beta2 * v + (1 - beta2) * g ** 2
                m_hat = m / (1 - beta1 ** t)
                v_hat = v / (1 - beta2 ** t)
                w -= lr * m_hat / (np.sqrt(v_hat) + eps)
            else:  # M+Adam: split mantissa/exponent optimization
                # Mantissa part: high-precision momentum
                m = beta1 * m + (1 - beta1) * g
                # Exponent part: sign-based update (low precision)
                exp_update = np.sign(g)
                v = beta2 * v + (1 - beta2) * g ** 2
                m_hat = m / (1 - beta1 ** t)
                v_hat = v / (1 - beta2 ** t)
                # Blend: 70% mantissa, 30% exponent
                update = 0.7 * m_hat / (np.sqrt(v_hat) + eps) + 0.3 * lr * exp_update
                w -= lr * update
            losses.append(loss_fn(w))
        return losses

    steps = np.arange(1, n_steps + 1)
    adam_losses = simulate("adam")
    madam_losses = simulate("madam")

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=steps, y=adam_losses, name="Adam (표준)", line=dict(color="#2196F3")))
    fig.add_trace(go.Scatter(x=steps, y=madam_losses, name="M+Adam (가수·지수 분리)", line=dict(color="#FF5722")))
    fig.update_layout(
        title="M+Adam vs Adam — 손실 곡선 비교",
        xaxis_title="Step", yaxis_title="Loss",
        template="plotly_dark", height=450, yaxis_type="log"
    )
    return fig


def build_madam_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 M+Adam: Low-Precision Training via Mantissa–Exponent Optimization
**ICML 2026 Spotlight** | Liang, Loeschcke, Anandkumar et al.

> 숫자의 **가수부(mantissa)**와 **지수부(exponent)**를 분리하여 각각 다르게 최적화 →
> 저정밀도(FP8/INT8)에서도 Adam 수준의 안정적 학습 달성.
> AI 훈련 비용 대폭 절감 가능.

**체험**: 파라미터를 조절해 M+Adam이 Adam 대비 어떻게 수렴하는지 비교하세요.
""")
        with gr.Row():
            lr = gr.Slider(0.001, 0.1, value=0.01, step=0.001, label="학습률 (lr)")
            beta1 = gr.Slider(0.5, 0.99, value=0.9, step=0.01, label="모멘텀 β₁")
            noise = gr.Slider(0.0, 2.0, value=0.5, step=0.1, label="그래디언트 노이즈")
            n_steps = gr.Slider(50, 500, value=200, step=50, label="Step 수")
        btn = gr.Button("비교 실행", variant="primary")
        plot = gr.Plot()
        btn.click(madam_simulation, inputs=[lr, beta1, noise, n_steps], outputs=plot)


# ── Tab 4: Linear RNN Parallelism ─────────────────────────────────────────────

def rnn_parallel_viz(seq_len: int, n_layers: int):
    seq_lens = np.arange(1, seq_len + 1)

    # Standard RNN: O(T) sequential — each step depends on previous
    rnn_latency = seq_lens * n_layers  # strictly sequential

    # Linear RNN: can use parallel scan (O(log T) depth)
    linear_rnn_latency = np.ceil(np.log2(np.maximum(seq_lens, 1))) * n_layers

    # Transformer: O(T²) attention but fully parallel — single pass
    transformer_latency = np.ones_like(seq_lens) * n_layers  # parallel, O(1) depth

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=seq_lens, y=rnn_latency, name="표준 RNN (순차, O(T))", line=dict(color="#F44336")))
    fig.add_trace(go.Scatter(x=seq_lens, y=linear_rnn_latency, name="선형 RNN (병렬 스캔, O(log T))", line=dict(color="#4CAF50")))
    fig.add_trace(go.Scatter(x=seq_lens, y=transformer_latency, name="Transformer (완전 병렬, O(1))", line=dict(color="#2196F3", dash="dot")))

    fig.update_layout(
        title=f"RNN 유형별 병렬화 깊이 비교 (레이어 {n_layers}개)",
        xaxis_title="시퀀스 길이 T", yaxis_title="병렬 처리 깊이 (낮을수록 빠름)",
        template="plotly_dark", height=450
    )
    return fig


def build_rnn_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 Why Are Linear RNNs More Parallelizable?
**ICML 2026 Poster** | Merrill, Jiang, Li, Lin, Sabharwal

> 표준 RNN은 각 스텝이 이전 스텝에 **의존** → 순차 실행 강제 → 병렬화 불가.
> 선형 RNN은 **병렬 스캔(prefix sum)** 알고리즘 적용 가능 → O(log T) 깊이로 단축.
> 이 논문은 이를 복잡도 클래스 이론으로 형식화.

**체험**: 시퀀스 길이가 늘어날수록 세 모델의 병렬 처리 깊이가 어떻게 달라지는지 확인하세요.
""")
        with gr.Row():
            seq_len = gr.Slider(10, 512, value=128, step=16, label="시퀀스 길이 T")
            n_layers = gr.Slider(1, 24, value=6, step=1, label="레이어 수")
        btn = gr.Button("비교 실행", variant="primary")
        plot = gr.Plot()
        btn.click(rnn_parallel_viz, inputs=[seq_len, n_layers], outputs=plot)


# ── Tab 5: Random Matrix Diffusion ────────────────────────────────────────────

def diffusion_consistency_viz(n_steps: int, noise_sigma: float, seed: int):
    np.random.seed(seed)
    img_size = 32
    # Simulate forward diffusion (noise injection)
    clean = np.random.randn(img_size, img_size) * 0.3 + 0.5
    clean = np.clip(clean, 0, 1)

    alphas = np.linspace(1.0, 0.0, n_steps + 1)  # noise schedule
    noisy_imgs = []
    for t in range(n_steps + 1):
        alpha = alphas[t]
        noise = np.random.randn(img_size, img_size) * noise_sigma
        noisy = alpha * clean + (1 - alpha) * noise
        noisy_imgs.append(noisy)

    # Plot selected timesteps
    show_steps = [0, n_steps // 4, n_steps // 2, 3 * n_steps // 4, n_steps]
    fig, _ = __import__('matplotlib.pyplot', fromlist=['']), None
    import matplotlib.pyplot as plt
    fig_mpl, axes = plt.subplots(1, len(show_steps), figsize=(14, 3))
    fig_mpl.patch.set_facecolor('#1a1a2e')
    for ax, t in zip(axes, show_steps):
        ax.imshow(noisy_imgs[t], cmap='viridis', vmin=-1, vmax=2)
        ax.set_title(f't={t}', color='white', fontsize=10)
        ax.axis('off')
    fig_mpl.suptitle("확산 과정 시각화 — 노이즈 주입 (t=0: 원본 → t=T: 순수 노이즈)", color='white')
    fig_mpl.tight_layout()
    return fig_mpl


def build_diffusion_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 A Random Matrix Perspective on the Consistency of Diffusion Models
**ICML 2026 Honorable Mention** | Wang, Zavatone-Veth, Pehlevan (Harvard)

> 확산 모델이 왜 비슷한 입력에 비슷한 결과를 내는지(일관성) **랜덤 행렬 이론**으로 수학적 설명.
> 무작위 가중치 행렬의 스펙트럼 특성 → 확산 궤적의 예측 가능성을 이론적으로 규명.

**체험**: 노이즈 주입 과정을 직접 시각화. 원본 이미지가 단계적으로 노이즈로 변환됨.
""")
        with gr.Row():
            n_steps = gr.Slider(10, 100, value=40, step=10, label="확산 단계 수 T")
            noise_sigma = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="노이즈 강도 σ")
            seed = gr.Slider(0, 99, value=42, step=1, label="랜덤 시드")
        btn = gr.Button("시각화", variant="primary")
        plot = gr.Plot(label="확산 과정")
        btn.click(diffusion_consistency_viz, inputs=[n_steps, noise_sigma, seed], outputs=plot)


# ── Tab 6: LM Circuits (ZeroGPU) ─────────────────────────────────────────────

@spaces.GPU
def lm_circuits_demo(text: str, layer_idx: int):
    from transformers import GPT2Model, GPT2Tokenizer
    import torch

    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = GPT2Model.from_pretrained("gpt2", output_hidden_states=True)
    model.eval()

    inputs = tokenizer(text, return_tensors="pt", max_length=32, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)

    # hidden state at chosen layer: [1, seq_len, 768]
    hidden = outputs.hidden_states[layer_idx][0].cpu().numpy()  # [seq_len, 768]
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

    # Sparsity: fraction of near-zero activations (|h| < threshold)
    threshold = 0.1
    sparsity_per_neuron = (np.abs(hidden) < threshold).mean(axis=0)  # [768]
    sparsity_per_token = (np.abs(hidden) < threshold).mean(axis=1)   # [seq_len]

    fig = go.Figure()
    # Heatmap: tokens × first 64 neurons
    n_show = min(64, hidden.shape[1])
    fig = go.Figure(data=go.Heatmap(
        z=hidden[:, :n_show],
        x=[f"N{i}" for i in range(n_show)],
        y=tokens,
        colorscale="RdBu",
        zmid=0,
    ))
    fig.update_layout(
        title=f"GPT-2 레이어 {layer_idx} — 뉴런 활성화 히트맵 (첫 64개 뉴런)\n"
              f"전체 뉴런 희소성: {sparsity_per_neuron.mean():.1%}",
        xaxis_title="뉴런 인덱스", yaxis_title="토큰",
        template="plotly_dark", height=max(300, len(tokens) * 30 + 100)
    )
    info = f"레이어 {layer_idx} 희소성: {sparsity_per_neuron.mean():.1%} (|활성화|<{threshold} 비율)\n"
    info += "희소성이 높을수록 → 회로 해석이 쉬움 (논문 핵심 주장)"
    return fig, info


def build_circuits_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 Language Model Circuits Are Sparse in the Neuron Basis
**ICML 2026 Spotlight** | Arora, Wu, Steinhardt, Schwettmann

> LLM 내부의 "회로"(특정 기능을 수행하는 뉴런 집합)가 **뉴런 기저에서 희소**함을 실증.
> 대부분 뉴런은 특정 시점에 거의 활성화되지 않음 → AI 블랙박스 해석 가능성 향상.

**체험**: 텍스트 입력 → GPT-2 각 레이어의 뉴런 활성화 패턴 시각화. 희소할수록 회로 해석 용이.

*ZeroGPU 사용 — 첫 실행 시 모델 로딩에 15–30초 소요될 수 있습니다.*
""")
        text_in = gr.Textbox(value="The cat sat on the mat", label="분석할 텍스트 (영어)")
        layer_idx = gr.Slider(0, 11, value=6, step=1, label="GPT-2 레이어 인덱스 (0~11)")
        btn = gr.Button("활성화 분석 (ZeroGPU)", variant="primary")
        plot = gr.Plot()
        info_out = gr.Textbox(label="희소성 정보", interactive=False)
        btn.click(lm_circuits_demo, inputs=[text_in, layer_idx], outputs=[plot, info_out])


# ── Tab 7: Any-Diffusion / SDXL Turbo (ZeroGPU) ──────────────────────────────

@spaces.GPU
def any_diffusion_demo(prompt: str, n_steps: int, seed: int):
    from diffusers import AutoPipelineForText2Image
    import torch

    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=torch.float16,
        variant="fp16"
    ).to("cuda")

    generator = torch.Generator("cuda").manual_seed(seed)
    image = pipe(
        prompt=prompt,
        num_inference_steps=n_steps,
        guidance_scale=0.0,
        generator=generator,
    ).images[0]
    return image


def build_anydiff_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 Any-Diffusion: Unified Multimodal Understanding and Generation
**ICML 2026 Poster** | Li, Long, Shen, Gao, Cao et al.

> 텍스트·이미지·음성을 **any-to-any** 방식으로 통합 처리하는 확산 모델.
> 마스크 기반 이산 확산으로 모달리티 경계를 제거.

**체험**: 프롬프트 입력 → 이미지 생성 (SDXL-Turbo, Any-Diffusion 개념 구현체).
SDXL-Turbo는 단 1–4 step으로 고품질 이미지 생성 — Any-Diffusion의 효율성 개념과 동일선상.

*ZeroGPU 사용 — 첫 실행 시 모델 로딩에 20–40초 소요될 수 있습니다.*
""")
        prompt = gr.Textbox(value="a futuristic AI laboratory with neural network visualizations", label="이미지 프롬프트 (영어)")
        with gr.Row():
            n_steps = gr.Slider(1, 4, value=2, step=1, label="추론 스텝 수 (1–4)")
            seed = gr.Slider(0, 9999, value=42, step=1, label="시드")
        btn = gr.Button("이미지 생성 (ZeroGPU)", variant="primary")
        img_out = gr.Image(label="생성 이미지")
        btn.click(any_diffusion_demo, inputs=[prompt, n_steps, seed], outputs=img_out)


# ── Tab 8: Flexibility Trap (LLM API comparison) ─────────────────────────────

@spaces.GPU
def flexibility_trap_demo(prompt: str, max_new_tokens: int):
    """Compare autoregressive (fixed order) vs masked diffusion LM (random order) generation."""
    from transformers import GPT2LMHeadModel, GPT2Tokenizer
    import torch

    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("gpt2").to("cuda")
    model.eval()

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = inputs["input_ids"].shape[1]

    # 1. Standard autoregressive (left-to-right, fixed order)
    with torch.no_grad():
        out_ar = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.8,
            pad_token_id=tokenizer.eos_token_id,
        )
    text_ar = tokenizer.decode(out_ar[0][input_len:], skip_special_tokens=True)

    # 2. Simulate "arbitrary order" by shuffling generation positions
    # (simplified proxy — actual diffusion LM requires MDLM architecture)
    with torch.no_grad():
        out_perm1 = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=1.2,  # higher temp = more "random" order effect
            pad_token_id=tokenizer.eos_token_id,
        )
    text_perm1 = tokenizer.decode(out_perm1[0][input_len:], skip_special_tokens=True)

    with torch.no_grad():
        out_perm2 = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=1.2,
            pad_token_id=tokenizer.eos_token_id,
        )
    text_perm2 = tokenizer.decode(out_perm2[0][input_len:], skip_special_tokens=True)

    result = f"""[고정 순서 생성 (Autoregressive — 왼쪽→오른쪽)]
{text_ar}

[임의 순서 시뮬레이션 #1 (higher temperature)]
{text_perm1}

[임의 순서 시뮬레이션 #2]
{text_perm2}

─────────────────────────────────────────────
논문 핵심: 임의 순서 생성은 표면적으로 다양해 보이지만
실제 추론 경로의 다양성은 오히려 낮아집니다 (The Flexibility Trap).
고정 순서 생성 #1이 대부분의 벤치마크에서 더 일관된 추론을 보입니다."""
    return result


def build_flexibility_tab():
    with gr.Column():
        gr.Markdown("""
## 📄 The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion LMs
**ICML 2026 Outstanding Paper** | Ni, Wang, Yue et al.

> 확산 언어 모델의 **임의 순서 생성**이 추론 다양성을 오히려 낮춘다는 역설적 발견.
> "유연할수록 더 획일적" — JustGRPO 훈련으로 개선.

**체험**: 같은 프롬프트로 고정 순서 vs 임의 순서 생성 결과를 비교하세요.
(GPT-2로 개념 시연 — 실제 확산 LM 행동의 proxy)

*ZeroGPU 사용*
""")
        prompt = gr.Textbox(value="The future of artificial intelligence will", label="시작 프롬프트")
        max_tokens = gr.Slider(20, 100, value=40, step=10, label="생성 토큰 수")
        btn = gr.Button("생성 비교 (ZeroGPU)", variant="primary")
        out = gr.Textbox(label="생성 결과 비교", lines=12)
        btn.click(flexibility_trap_demo, inputs=[prompt, max_tokens], outputs=out)


# ── Main App ──────────────────────────────────────────────────────────────────

with gr.Blocks(
    title="AI Paper Lab — ICML 2026",
    theme=gr.themes.Base(primary_hue="violet"),
    css="""
    .paper-header { background: linear-gradient(135deg, #1a1a2e, #16213e); padding: 20px; border-radius: 10px; margin-bottom: 10px; }
    """
) as demo:
    gr.Markdown("""
# AI Paper Lab
### ICML 2026 주요 논문을 직접 체험하세요

각 탭에서 논문의 핵심 개념을 인터랙티브하게 경험할 수 있습니다.
CPU 탭은 즉시 실행, ZeroGPU 탭은 첫 실행 시 20–40초 대기 발생.
""")

    with gr.Tabs():
        with gr.Tab("1. Grok Grokking (CPU)"):
            build_grokking_tab()
        with gr.Tab("2. Quantum vs Classical (CPU)"):
            build_quantum_tab()
        with gr.Tab("3. M+Adam (CPU)"):
            build_madam_tab()
        with gr.Tab("4. Linear RNN (CPU)"):
            build_rnn_tab()
        with gr.Tab("5. Diffusion 시각화 (CPU)"):
            build_diffusion_tab()
        with gr.Tab("6. LM Circuits (ZeroGPU)"):
            build_circuits_tab()
        with gr.Tab("7. Any-Diffusion (ZeroGPU)"):
            build_anydiff_tab()
        with gr.Tab("8. Flexibility Trap (ZeroGPU)"):
            build_flexibility_tab()

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