ui/gradio_app.py

"""
Gradio web UI for Indic Heritage Studio v2.

6 tabs:
  1. Text → Image (SDXL + per-style LoRA)
  2. Style Transfer (IP-Adapter XL)
  3. Image → Video (Stable Video Diffusion)
  4. ControlNet (composition-conditioned generation)
  5. Inpainting (mask + restyle)
  6. Batch Processor (multi-GPU parallel)

Plus: Style Advisor widget (free AMD Qwen API), GPU monitor sidebar.
"""
from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import List, Optional

import gradio as gr
from PIL import Image

from config.settings import settings
from config.styles import HERITAGE_STYLES, StyleSpec, get_style, list_styles
from agents.style_advisor import StyleAdvisor
from agents.critic import Critic
from utils.gpu_utils import list_gpus

log = logging.getLogger(__name__)

STYLE_CHOICES = [s.id for s in list_styles()]
STYLE_LABELS = {s.id: f"{s.display_name} ({s.region})" for s in list_styles()}


def _style_dropdown(label: str = "Heritage Style"):
    return gr.Dropdown(
        choices=STYLE_CHOICES,
        value="madhubani",
        label=label,
        show_label=True,
        type="value",
        allow_custom_value=False,
    )


def _gpu_status_html() -> str:
    gpus = list_gpus()
    if not gpus:
        return "<p><b>GPU:</b> CPU only (CUDA/ROCm not detected)</p>"
    rows = "".join(
        f"<tr><td>GPU {g.index}</td><td>{g.name}</td>"
        f"<td>{g.vram_total_gb:.1f} GB</td>"
        f"<td>{g.vram_free_gb:.1f} GB free</td></tr>"
        for g in gpus
    )
    return f"""
    <table style="width:100%; font-size: 12px; border-collapse: collapse;">
      <thead>
        <tr><th>#</th><th>Name</th><th>VRAM</th><th>Free</th></tr>
      </thead>
      <tbody>{rows}</tbody>
    </table>
    <p><b>Total VRAM:</b> {settings.total_vram_gb:.1f} GB across {len(gpus)} GPUs</p>
    """


# ---------------------------------------------------------------------------
# Tab 1: Text → Image
# ---------------------------------------------------------------------------
_t2i_pipe = None


def _get_t2i():
    global _t2i_pipe
    if _t2i_pipe is None:
        from core.text_to_image import TextToImagePipeline
        _t2i_pipe = TextToImagePipeline().load()
    return _t2i_pipe


def t2i_generate(prompt, style_id, steps, guidance, size, seed, use_lora, high_quality):
    pipe = _get_t2i()
    style = get_style(style_id)
    img = pipe.generate(
        prompt=prompt, style=style,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(size), height=int(size),
        seed=int(seed),
        use_lora=bool(use_lora),
        high_quality=bool(high_quality),
    )
    # Critique
    critique = Critic().evaluate(img, style, prompt)
    info = (
        f"Style: {style.display_name} | "
        f"Steps: {steps} | Size: {size}×{size} | "
        f"LoRA: {'on' if use_lora else 'off'} | "
        f"Quality: {'high' if high_quality else 'standard'}\n"
        f"Critic: {critique.overall:.1f}/10 "
        f"(fidelity={critique.style_fidelity}, comp={critique.composition}, "
        f"tech={critique.technical_quality}) — {critique.feedback}"
    )
    return img, info


# ---------------------------------------------------------------------------
# Tab 2: Style Transfer
# ---------------------------------------------------------------------------
_style_pipe = None


def _get_style_pipe():
    global _style_pipe
    if _style_pipe is None:
        from core.style_transfer import StyleTransferPipeline
        _style_pipe = StyleTransferPipeline().load()
    return _style_pipe


def style_transfer_run(image, style_id, strength, ip_scale, steps, seed, use_lora):
    if image is None:
        raise gr.Error("Please upload an input image")
    pipe = _get_style_pipe()
    style = get_style(style_id)
    out = pipe.transfer(
        image=image, style=style,
        strength=float(strength),
        ip_adapter_scale=float(ip_scale),
        num_inference_steps=int(steps),
        seed=int(seed),
        use_lora=bool(use_lora),
    )
    return out, f"Transferred {style.display_name} style. Strength={strength}, IP scale={ip_scale}"


# ---------------------------------------------------------------------------
# Tab 3: Image → Video
# ---------------------------------------------------------------------------
_i2v_pipe = None


def _get_i2v():
    global _i2v_pipe
    if _i2v_pipe is None:
        from core.image_to_video import ImageToVideoPipeline
        _i2v_pipe = ImageToVideoPipeline().load()
    return _i2v_pipe


def i2v_run(image, style_id, num_frames, fps, motion_bucket, seed, width, height):
    if image is None:
        raise gr.Error("Please upload an input image")
    pipe = _get_i2v()
    style = get_style(style_id) if style_id else None
    frames = pipe.generate(
        image=image, style=style,
        num_frames=int(num_frames), fps=int(fps),
        motion_bucket_id=int(motion_bucket) if motion_bucket else None,
        seed=int(seed),
        width=int(width), height=int(height),
    )
    # Save as GIF for Gradio preview + return mp4 path
    from utils.video_utils import gif_from_frames, frames_to_mp4
    out_dir = settings.outputs_dir / "ui" / "i2v"
    out_dir.mkdir(parents=True, exist_ok=True)
    ts = int(time.time())
    mp4_path = out_dir / f"i2v_{ts}.mp4"
    gif_path = out_dir / f"i2v_{ts}.gif"
    frames_to_mp4(frames, mp4_path, fps=int(fps))
    gif_from_frames(frames, gif_path, fps=int(fps))
    return str(mp4_path), str(gif_path), f"Generated {len(frames)} frames @ {fps} fps"


# ---------------------------------------------------------------------------
# Tab 4: ControlNet
# ---------------------------------------------------------------------------
_controlnet_pipes = {}


def _get_controlnet(condition_type):
    if condition_type not in _controlnet_pipes:
        from core.controlnet import ControlNetPipeline
        _controlnet_pipes[condition_type] = ControlNetPipeline(
            condition_type=condition_type
        ).load()
    return _controlnet_pipes[condition_type]


def controlnet_run(condition_image, condition_type, prompt, style_id, scale, steps, seed, size):
    if condition_image is None:
        raise gr.Error("Please upload a conditioning image")
    pipe = _get_controlnet(condition_type)
    style = get_style(style_id)

    # 1. Detect
    cond = pipe.detect(condition_image)
    # 2. Generate
    out = pipe.generate(
        conditioning_image=cond,
        prompt=prompt, style=style,
        controlnet_conditioning_scale=float(scale),
        num_inference_steps=int(steps),
        width=int(size), height=int(size),
        seed=int(seed),
    )
    return cond, out, f"ControlNet ({condition_type}) → {style.display_name}"


# ---------------------------------------------------------------------------
# Tab 5: Inpainting
# ---------------------------------------------------------------------------
_inpaint_pipe = None


def _get_inpaint():
    global _inpaint_pipe
    if _inpaint_pipe is None:
        from core.inpainting import InpaintingPipeline
        _inpaint_pipe = InpaintingPipeline().load()
    return _inpaint_pipe


def inpaint_run(image, mask, prompt, style_id, steps, strength, seed, use_lora):
    if image is None or mask is None:
        raise gr.Error("Please provide both an image and a mask")
    pipe = _get_inpaint()
    style = get_style(style_id)
    out = pipe.inpaint(
        image=image, mask=mask,
        prompt=prompt, style=style,
        num_inference_steps=int(steps),
        strength=float(strength),
        seed=int(seed),
        use_lora=bool(use_lora),
    )
    return out, f"Inpainted region in {style.display_name} style"


# ---------------------------------------------------------------------------
# Tab 6: Batch Processor
# ---------------------------------------------------------------------------
def batch_run(prompt, mode, input_dir, styles, seeds, num_workers, output_dir):
    from core.batch_processor import BatchProcessor, BatchJob
    import asyncio

    if mode == "t2i":
        jobs = BatchProcessor.build_t2i_jobs(
            prompt=prompt, styles=styles,
            seeds=[int(s) for s in seeds.split(",")] if seeds else [42, 1337, 2024],
            output_dir=Path(output_dir),
        )
    elif mode == "style_transfer":
        if not input_dir:
            raise gr.Error("--input-dir required for style_transfer mode")
        jobs = BatchProcessor.build_style_transfer_jobs(
            input_dir=Path(input_dir), styles=styles,
            output_dir=Path(output_dir),
        )
    else:
        raise gr.Error(f"Unknown mode: {mode}")

    processor = BatchProcessor(num_workers=int(num_workers) if num_workers else None)
    results = processor.run(jobs)

    succ = sum(1 for r in results if r.success)
    fail = len(results) - succ
    total_time = sum(r.elapsed_seconds for r in results)
    return (
        f"✅ {succ} succeeded, ❌ {fail} failed\n"
        f"Total worker time: {total_time:.1f}s\n"
        f"Effective throughput: {succ / max(total_time, 1) * 60:.1f} images/min\n"
        f"Output dir: {output_dir}"
    )


# ---------------------------------------------------------------------------
# Style Advisor
# ---------------------------------------------------------------------------
_advisor = None


def _get_advisor():
    global _advisor
    if _advisor is None:
        _advisor = StyleAdvisor()
    return _advisor


def advisor_recommend(prompt):
    adv = _get_advisor()
    result = adv.recommend(prompt)
    style = get_style(result["style"])
    text = (
        f"### Recommended: {style.display_name}\n\n"
        f"**Region:** {style.region}\n\n"
        f"**Reason:** {result.get('reason', 'No rationale provided')}\n\n"
        f"**Confidence:** {result.get('confidence', 0):.0%}\n\n"
        f"**Source:** `{result.get('source', 'unknown')}`\n\n"
        f"**Description:** {style.description}\n\n"
        f"**Cultural keywords:** {', '.join(style.cultural_keywords)}"
    )
    return result["style"], text


# ---------------------------------------------------------------------------
# Build UI
# ---------------------------------------------------------------------------
def build_ui():
    with gr.Blocks(
        title="Indic Heritage Studio v2",
        theme=gr.themes.Soft(primary_hue="amber", secondary_hue="blue"),
        css="""
        .gpu-monitor { font-family: monospace; background: #1f2937; color: #e5e7eb;
                       padding: 12px; border-radius: 8px; font-size: 12px; }
        .style-card { border-left: 4px solid #f59e0b; padding-left: 12px; margin: 8px 0; }
        """,
    ) as demo:
        gr.Markdown("# 🎨 Indic Heritage Studio v2")
        gr.Markdown(
            "Multimodal content creation tool that reimagines modern photos and prompts "
            "through the lens of Indian heritage art forms — running on SDXL + SVD + "
            "ControlNet + per-style LoRAs across 8 × NVIDIA 80GB GPUs."
        )

        with gr.Accordion("🖥 GPU Monitor", open=False):
            gpu_html = gr.HTML(value=_gpu_status_html())
            refresh_btn = gr.Button("Refresh", size="sm")
            refresh_btn.click(fn=_gpu_status_html, outputs=gpu_html)

        # Style Advisor (always visible at top)
        with gr.Accordion("🧭 AI Style Advisor (free AMD Qwen API)", open=False):
            with gr.Row():
                advisor_prompt = gr.Textbox(
                    label="Describe what you want to create",
                    placeholder="e.g. a young woman reading under a banyan tree at sunset",
                    lines=2,
                )
                advisor_btn = gr.Button("Get Recommendation", variant="primary")
            with gr.Row():
                advisor_style = gr.Dropdown(
                    choices=STYLE_CHOICES, value="madhubani", label="Suggested style",
                )
                advisor_explain = gr.Markdown()

            advisor_btn.click(
                fn=advisor_recommend,
                inputs=advisor_prompt,
                outputs=[advisor_style, advisor_explain],
            )

        # Tab 1: Text → Image
        with gr.Tab("🖼 Text → Image (SDXL)"):
            with gr.Row():
                with gr.Column(scale=1):
                    t2i_prompt = gr.Textbox(
                        label="Prompt", lines=3,
                        value="a young woman reading under a banyan tree at sunset",
                    )
                    t2i_style = _style_dropdown()
                    with gr.Row():
                        t2i_steps = gr.Slider(10, 80, value=25, step=5, label="Steps")
                        t2i_guidance = gr.Slider(1, 15, value=7.0, step=0.5, label="CFG")
                    with gr.Row():
                        t2i_size = gr.Slider(512, 1536, value=1024, step=64,
                                             label="Resolution (square)")
                        t2i_seed = gr.Number(value=42, label="Seed", precision=0)
                    with gr.Row():
                        t2i_lora = gr.Checkbox(value=True, label="Use style LoRA")
                        t2i_hq = gr.Checkbox(value=False, label="High-quality (50 steps + refiner)")
                    t2i_btn = gr.Button("Generate", variant="primary")
                with gr.Column(scale=1):
                    t2i_out = gr.Image(label="Generated Image", type="pil")
                    t2i_info = gr.Textbox(label="Metadata + Critic", lines=5)
            t2i_btn.click(
                fn=t2i_generate,
                inputs=[t2i_prompt, t2i_style, t2i_steps, t2i_guidance,
                        t2i_size, t2i_seed, t2i_lora, t2i_hq],
                outputs=[t2i_out, t2i_info],
            )

        # Tab 2: Style Transfer
        with gr.Tab("🎨 Style Transfer (IP-Adapter XL)"):
            with gr.Row():
                with gr.Column(scale=1):
                    st_image = gr.Image(label="Input Image", type="pil")
                    st_style = _style_dropdown()
                    with gr.Row():
                        st_strength = gr.Slider(0.0, 1.0, value=0.7, step=0.05,
                                                label="Strength (denoising)")
                        st_ipscale = gr.Slider(0.0, 1.0, value=0.7, step=0.05,
                                               label="IP-Adapter scale")
                    with gr.Row():
                        st_steps = gr.Slider(10, 80, value=30, step=5, label="Steps")
                        st_seed = gr.Number(value=42, label="Seed", precision=0)
                    st_lora = gr.Checkbox(value=True, label="Use style LoRA overlay")
                    st_btn = gr.Button("Transfer Style", variant="primary")
                with gr.Column(scale=1):
                    st_out = gr.Image(label="Stylized Image", type="pil")
                    st_info = gr.Textbox(label="Info", lines=2)
            st_btn.click(
                fn=style_transfer_run,
                inputs=[st_image, st_style, st_strength, st_ipscale,
                        st_steps, st_seed, st_lora],
                outputs=[st_out, st_info],
            )

        # Tab 3: Image → Video
        with gr.Tab("🎬 Image → Video (SVD)"):
            with gr.Row():
                with gr.Column(scale=1):
                    iv_image = gr.Image(label="Input Image (still)", type="pil")
                    iv_style = _style_dropdown("Heritage Style (motion hint)")
                    with gr.Row():
                        iv_frames = gr.Slider(8, 50, value=25, step=1, label="Frames")
                        iv_fps = gr.Slider(4, 16, value=8, step=1, label="FPS")
                    with gr.Row():
                        iv_motion = gr.Slider(1, 255, value=127, step=1,
                                              label="Motion bucket (1=subtle, 255=dynamic)")
                        iv_seed = gr.Number(value=42, label="Seed", precision=0)
                    with gr.Row():
                        iv_w = gr.Slider(512, 1280, value=1024, step=64, label="Width")
                        iv_h = gr.Slider(320, 768, value=576, step=64, label="Height")
                    iv_btn = gr.Button("Generate Video", variant="primary")
                with gr.Column(scale=1):
                    iv_mp4 = gr.Video(label="Output MP4")
                    iv_gif = gr.Image(label="Preview GIF")
                    iv_info = gr.Textbox(label="Info", lines=2)
            iv_btn.click(
                fn=i2v_run,
                inputs=[iv_image, iv_style, iv_frames, iv_fps, iv_motion,
                        iv_seed, iv_w, iv_h],
                outputs=[iv_mp4, iv_gif, iv_info],
            )

        # Tab 4: ControlNet
        with gr.Tab("📐 ControlNet"):
            with gr.Row():
                with gr.Column(scale=1):
                    cn_image = gr.Image(label="Conditioning Image", type="pil")
                    cn_type = gr.Radio(
                        choices=["canny", "depth", "openpose"],
                        value="canny", label="Condition type",
                    )
                    cn_prompt = gr.Textbox(
                        label="Prompt", lines=2,
                        value="a courtly gathering with musicians and dancers",
                    )
                    cn_style = _style_dropdown()
                    with gr.Row():
                        cn_scale = gr.Slider(0.0, 2.0, value=0.8, step=0.1,
                                             label="ControlNet scale")
                        cn_steps = gr.Slider(10, 80, value=30, step=5, label="Steps")
                    with gr.Row():
                        cn_seed = gr.Number(value=42, label="Seed", precision=0)
                        cn_size = gr.Slider(512, 1536, value=1024, step=64,
                                            label="Resolution")
                    cn_btn = gr.Button("Generate with ControlNet", variant="primary")
                with gr.Column(scale=1):
                    cn_cond = gr.Image(label="Detected Condition", type="pil")
                    cn_out = gr.Image(label="Generated Heritage Art", type="pil")
                    cn_info = gr.Textbox(label="Info", lines=2)
            cn_btn.click(
                fn=controlnet_run,
                inputs=[cn_image, cn_type, cn_prompt, cn_style,
                        cn_scale, cn_steps, cn_seed, cn_size],
                outputs=[cn_cond, cn_out, cn_info],
            )

        # Tab 5: Inpainting
        with gr.Tab("✏ Inpainting"):
            with gr.Row():
                with gr.Column(scale=1):
                    ip_image = gr.Image(label="Input Image", type="pil")
                    ip_mask = gr.Image(label="Mask (white = inpaint)", type="pil",
                                       image_mode="L")
                    ip_prompt = gr.Textbox(
                        label="What to fill", lines=2,
                        value="ornate floral border with peacock motifs",
                    )
                    ip_style = _style_dropdown()
                    with gr.Row():
                        ip_steps = gr.Slider(10, 80, value=30, step=5, label="Steps")
                        ip_strength = gr.Slider(0.0, 1.0, value=1.0, step=0.05,
                                                label="Strength")
                    with gr.Row():
                        ip_seed = gr.Number(value=42, label="Seed", precision=0)
                        ip_lora = gr.Checkbox(value=True, label="Use style LoRA")
                    ip_btn = gr.Button("Inpaint", variant="primary")
                with gr.Column(scale=1):
                    ip_out = gr.Image(label="Inpainted Result", type="pil")
                    ip_info = gr.Textbox(label="Info", lines=2)
            ip_btn.click(
                fn=inpaint_run,
                inputs=[ip_image, ip_mask, ip_prompt, ip_style,
                        ip_steps, ip_strength, ip_seed, ip_lora],
                outputs=[ip_out, ip_info],
            )

        # Tab 6: Batch Processor
        with gr.Tab("⚡ Batch (Multi-GPU)"):
            with gr.Row():
                with gr.Column(scale=1):
                    bt_mode = gr.Radio(
                        choices=["t2i", "style_transfer"], value="t2i",
                        label="Batch mode",
                    )
                    bt_prompt = gr.Textbox(
                        label="Prompt (T2I mode)", lines=2,
                        value="a temple festival at dawn",
                    )
                    bt_input_dir = gr.Textbox(
                        label="Input dir (style_transfer mode)",
                        value="examples/inputs",
                    )
                    bt_styles = gr.CheckboxGroup(
                        choices=STYLE_CHOICES, value=STYLE_CHOICES,
                        label="Styles (one per image)",
                    )
                    with gr.Row():
                        bt_seeds = gr.Textbox(
                            value="42,1337,2024",
                            label="Seeds (comma-separated, T2I mode)",
                        )
                        bt_workers = gr.Slider(1, 8, value=4, step=1,
                                               label="Parallel GPU workers")
                    bt_output = gr.Textbox(
                        value="outputs/batch",
                        label="Output directory",
                    )
                    bt_btn = gr.Button("Run Batch", variant="primary")
                with gr.Column(scale=1):
                    bt_log = gr.Textbox(label="Batch Result", lines=8)
            bt_btn.click(
                fn=batch_run,
                inputs=[bt_prompt, bt_mode, bt_input_dir, bt_styles,
                        bt_seeds, bt_workers, bt_output],
                outputs=bt_log,
            )

        # Footer
        gr.Markdown("---")
        gr.Markdown(
            f"**Indic Heritage Studio v2** — running on `{settings.gpu_name}` "
            f"with {settings.device_count} GPU(s), {settings.total_vram_gb:.0f} GB total VRAM. "
            "Track 1 submission for AMD AI DevMaster Hackathon 2026."
        )

    return demo