"""RxBrain — Embodied Cognition Foundation Model demo (ZeroGPU).

tencent/Hy-Embodied-RxBrain-1.0 : a unified Mixture-of-Transformers model that
couples language reasoning with visual imagination. This Space exposes three
capabilities, each mapped 1:1 onto the authors' official inference scripts:

  * Visual Question Answering (vqa_inference.answer)
  * Text-to-Image imagination (text2image_inference.generate_image)
  * Interleaved embodied planning (interleave_inference.generate_step_joint)
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — must precede torch / CUDA-touching imports
import tempfile  # noqa: E402
import traceback  # noqa: E402

import torch  # noqa: E402
import numpy as np  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image  # noqa: E402
from huggingface_hub import snapshot_download, hf_hub_download  # noqa: E402

# ---------------------------------------------------------------------------
# Model / VAE loading (module scope, eager .to("cuda") for ZeroGPU)
# ---------------------------------------------------------------------------
MODEL_ID = "tencent/Hy-Embodied-RxBrain-1.0"
# Ungated mirror of the FLUX autoencoder (ae.safetensors) — byte-for-byte the
# same tensors (244 keys, identical shapes) as black-forest-labs/FLUX.1-schnell,
# but without the gate, so the Space loads it without an HF token.
VAE_REPO = "Comfy-Org/Lumina_Image_2.0_Repackaged"
VAE_FILE = "split_files/vae/ae.safetensors"
DTYPE = torch.bfloat16

print("Downloading RxBrain checkpoint …")
CKPT_DIR = snapshot_download(MODEL_ID)
print(f"Checkpoint at {CKPT_DIR}")

print("Downloading FLUX VAE (ae.safetensors) …")
VAE_PATH = hf_hub_download(VAE_REPO, VAE_FILE)
print(f"VAE at {VAE_PATH}")

from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor  # noqa: E402
from model import UnifiedMoTForConditionalGeneration, maybe_init_generation_path  # noqa: E402
from vae_model.autoencoder import load_ae  # noqa: E402

# Inference helpers ported straight from the official scripts.
from vqa_inference import answer as _vqa_answer  # noqa: E402
from text2image_inference import generate_image as _t2i_generate  # noqa: E402
from interleave_inference import generate_step_joint as _interleave_step  # noqa: E402

print("Loading processor …")
processor = HunYuanVLMoTProcessor.from_pretrained(CKPT_DIR, trust_remote_code=True)

print("Loading model …")
model = UnifiedMoTForConditionalGeneration.from_pretrained(CKPT_DIR, dtype=DTYPE)
maybe_init_generation_path(model, model_load_path=CKPT_DIR)
model.to("cuda").eval()

print("Loading VAE …")
vae, _ = load_ae(VAE_PATH)
vae.requires_grad_(False)
vae.eval()
vae.to("cuda", dtype=DTYPE)

DEVICE = torch.device("cuda")
print("All components ready.")


def _tensor_to_pil(img):
    """(3, H, W) tensor in [-1, 1] -> PIL.Image."""
    arr = ((img.cpu().permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8")
    return Image.fromarray(arr)


def _save_input_image(image):
    """Persist a Gradio image (PIL or path) to a temp jpg, return the path."""
    if image is None:
        return None
    if isinstance(image, str):
        return image
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    f = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    image.convert("RGB").save(f.name, quality=95)
    return f.name


# ---------------------------------------------------------------------------
# 1. Visual Question Answering
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def vqa(image, question, max_new_tokens=256):
    """Answer a question about an image using RxBrain's understanding path.

    Args:
        image: the input image to reason over.
        question: a natural-language question about the image.
        max_new_tokens: maximum number of answer tokens to decode.
    Returns:
        The model's answer text.
    """
    if image is None:
        return "Please provide an input image."
    if not question or not question.strip():
        return "Please enter a question."
    path = _save_input_image(image)
    text = _vqa_answer(
        model, processor, image_paths=[path], question=question.strip(),
        device=DEVICE, dtype=DTYPE, max_new_tokens=int(max_new_tokens),
    )
    # The model wraps answers in <answer>...</answer>; strip the tags for display.
    text = text.strip()
    for tag in ("<answer>", "</answer>"):
        text = text.replace(tag, "")
    return text.strip()


# ---------------------------------------------------------------------------
# 2. Text-to-Image imagination
# ---------------------------------------------------------------------------
def _t2i_duration(prompt, height=256, width=256, num_steps=25, cfg_scale=1.0, seed=0):
    base = 40
    per_step = 1.5 * (2 if cfg_scale != 1.0 else 1)
    return int(min(180, base + num_steps * per_step))


@spaces.GPU(duration=_t2i_duration)
def text_to_image(prompt, height=256, width=256, num_steps=25, cfg_scale=1.0, seed=0):
    """Imagine an image from a text prompt via the flow-matching head.

    Args:
        prompt: the text description to imagine.
        height: output image height in pixels (multiple of 16).
        width: output image width in pixels (multiple of 16).
        num_steps: number of Euler-ODE flow-matching steps.
        cfg_scale: classifier-free guidance scale (>1 enables CFG).
        seed: RNG seed for reproducibility.
    Returns:
        The generated PIL image.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    torch.manual_seed(int(seed))
    h = int(height) // 16 * 16
    w = int(width) // 16 * 16
    img = _t2i_generate(
        model, vae, processor, prompt.strip(),
        height=h, width=w, num_steps=int(num_steps),
        device=DEVICE, dtype=DTYPE, cfg_scale=float(cfg_scale),
    )
    return _tensor_to_pil(img)


# ---------------------------------------------------------------------------
# 3. Interleaved Embodied Planning (text plan + goal images, step by step)
# ---------------------------------------------------------------------------
def _plan_duration(image, task, num_frames=3, num_steps=50, height=144, width=256, seed=0):
    return int(min(230, 30 + num_frames * (25 + num_steps * 1.3)))


@spaces.GPU(duration=_plan_duration)
def embodied_plan(image, task, num_frames=3, num_steps=50, height=144, width=256, seed=0):
    """Decompose an embodied task into steps, emitting both the next action
    (text) and the imagined goal frame (image) for each step.

    Args:
        image: an observation frame of the current scene.
        task: the embodied task to plan for.
        num_frames: number of interleaved plan steps to roll out.
        num_steps: number of flow-matching ODE steps per imagined frame.
        height: imagined-frame height in pixels (multiple of 16).
        width: imagined-frame width in pixels (multiple of 16).
        seed: RNG seed for reproducibility.
    Yields:
        (gallery_of_goal_frames, plan_markdown) as each step completes.
    """
    if image is None:
        raise gr.Error("Please provide an observation image.")
    if not task or not task.strip():
        raise gr.Error("Please enter a task description.")
    torch.manual_seed(int(seed))
    obs_path = _save_input_image(image)
    h = int(height) // 16 * 16
    w = int(width) // 16 * 16

    prior_steps = []          # (text, frame_path) accumulated across steps
    gallery = []              # (image, caption) for the gr.Gallery
    plan_md_lines = [f"**Task:** {task.strip()}", ""]

    for k in range(int(num_frames)):
        text, img = _interleave_step(
            model, vae, processor,
            input_image_paths=[obs_path], task_text=task.strip(),
            height=h, width=w, num_steps=int(num_steps),
            device=DEVICE, dtype=DTYPE, max_text_tokens=96,
            prior_steps=list(prior_steps),
        )
        pil = _tensor_to_pil(img)
        frame_f = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        pil.save(frame_f.name)
        prior_steps.append((text, frame_f.name))
        gallery.append((pil, f"Step {k + 1}"))
        # Tidy the emitted text: drop any wrapper tags and a leading "Step N:"
        # the model may already produce (we add our own numbered heading).
        disp = text.strip()
        for tag in ("<answer>", "</answer>"):
            disp = disp.replace(tag, "")
        disp = disp.strip()
        if disp.lower().startswith(f"step {k + 1}:"):
            disp = disp[len(f"step {k + 1}:"):].strip()
        plan_md_lines.append(f"**Step {k + 1}:** {disp or '(imagined goal frame only)'}")
        yield gallery, "\n\n".join(plan_md_lines)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = """# 🔮 RxBrain — Embodied Cognition Foundation Model

**[tencent/Hy-Embodied-RxBrain-1.0](https://huggingface.co/tencent/Hy-Embodied-RxBrain-1.0)** — a
unified Mixture-of-Transformers model that couples **language reasoning** with **visual imagination**.
It answers questions about scenes, imagines images from text, and decomposes embodied tasks into
interleaved step-by-step *plans + goal images*.

*Tencent Robotics X × Futian Laboratory × Tencent Hunyuan · imagination decoded through a frozen FLUX VAE.*
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="RxBrain") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)

        with gr.Tabs():
            # ---------------- VQA ----------------
            with gr.Tab("🧠 Visual QA"):
                with gr.Row():
                    with gr.Column():
                        vqa_image = gr.Image(label="Image", type="pil", height=320)
                        vqa_question = gr.Textbox(
                            label="Question",
                            placeholder="What objects are on the stovetop?",
                        )
                        with gr.Accordion("Advanced", open=False):
                            vqa_max_tokens = gr.Slider(
                                16, 512, value=256, step=16, label="Max new tokens")
                        vqa_btn = gr.Button("Answer", variant="primary")
                    with gr.Column():
                        vqa_out = gr.Textbox(label="Answer", lines=10)

                vqa_btn.click(
                    vqa, inputs=[vqa_image, vqa_question, vqa_max_tokens],
                    outputs=vqa_out, api_name="vqa",
                )
                gr.Examples(
                    examples=[
                        ["examples/bridgev2_obs.jpg",
                         "What objects are on the stovetop, and where is the green toy?"],
                        ["examples/umi_sock_obs.jpg",
                         "Describe the scene and the objects the robot could manipulate."],
                        ["examples/xtrainer_garment_obs.jpg",
                         "What garment is on the surface and what is its current state?"],
                    ],
                    inputs=[vqa_image, vqa_question],
                    outputs=vqa_out, fn=vqa,
                    cache_examples=False, run_on_click=True,
                )

            # ---------------- T2I ----------------
            with gr.Tab("🎨 Text-to-Image"):
                with gr.Row():
                    with gr.Column():
                        t2i_prompt = gr.Textbox(
                            label="Prompt",
                            placeholder="a watercolor painting of a cat",
                        )
                        with gr.Accordion("Advanced", open=False):
                            with gr.Row():
                                t2i_h = gr.Slider(128, 384, value=256, step=16, label="Height")
                                t2i_w = gr.Slider(128, 384, value=256, step=16, label="Width")
                            t2i_steps = gr.Slider(10, 50, value=25, step=1, label="ODE steps")
                            t2i_cfg = gr.Slider(1.0, 7.0, value=1.0, step=0.5,
                                                label="CFG scale (>1 enables guidance)")
                            t2i_seed = gr.Number(value=0, precision=0, label="Seed")
                        t2i_btn = gr.Button("Imagine", variant="primary")
                    with gr.Column():
                        t2i_out = gr.Image(label="Imagined image", height=320)

                t2i_btn.click(
                    text_to_image,
                    inputs=[t2i_prompt, t2i_h, t2i_w, t2i_steps, t2i_cfg, t2i_seed],
                    outputs=t2i_out, api_name="text_to_image",
                )
                gr.Examples(
                    examples=[
                        ["a watercolor painting of a cat"],
                        ["a red apple on a wooden table, photorealistic"],
                        ["a robot arm gripping a colorful toy on a kitchen counter"],
                    ],
                    inputs=[t2i_prompt],
                    outputs=t2i_out, fn=text_to_image,
                    cache_examples=False, run_on_click=True,
                )

            # ---------------- Interleaved Planning ----------------
            with gr.Tab("🧩 Embodied Planning"):
                gr.Markdown(
                    "Give an observation frame + a task; RxBrain rolls out an interleaved "
                    "plan — each step emits an **action (text)** and the **imagined goal frame**."
                )
                with gr.Row():
                    with gr.Column():
                        plan_image = gr.Image(label="Observation", type="pil", height=280)
                        plan_task = gr.Textbox(label="Task", lines=3,
                                               placeholder="Move the green toy over the metal pot.")
                        with gr.Accordion("Advanced", open=False):
                            plan_nframes = gr.Slider(1, 5, value=3, step=1, label="Plan steps")
                            plan_steps = gr.Slider(10, 50, value=50, step=1, label="ODE steps/frame")
                            with gr.Row():
                                plan_h = gr.Slider(128, 256, value=144, step=16, label="Frame height")
                                plan_w = gr.Slider(128, 384, value=256, step=16, label="Frame width")
                            plan_seed = gr.Number(value=0, precision=0, label="Seed")
                        plan_btn = gr.Button("Plan", variant="primary")
                    with gr.Column():
                        plan_gallery = gr.Gallery(label="Imagined goal frames", columns=3, height=280)
                        plan_text = gr.Markdown(label="Plan")

                plan_btn.click(
                    embodied_plan,
                    inputs=[plan_image, plan_task, plan_nframes, plan_steps,
                            plan_h, plan_w, plan_seed],
                    outputs=[plan_gallery, plan_text], api_name="embodied_plan",
                )
                gr.Examples(
                    examples=[
                        ["examples/bridgev2_obs.jpg",
                         "Move the green toy from the left side of the stovetop and hold it over "
                         "the metal pot on the left burner. Please give the full plan as concrete "
                         "physical actions, one step per action, and imagine the state once each "
                         "step is completed. If only one step remains, provide only that final step."],
                        ["examples/umi_sock_obs.jpg",
                         "Fold the purple sock in half. Please lay out the high-level subgoal plan "
                         "and imagine the resulting frame after each major stage. If only one "
                         "subgoal remains, provide only that final subgoal."],
                    ],
                    inputs=[plan_image, plan_task],
                    outputs=[plan_gallery, plan_text], fn=embodied_plan,
                    cache_examples=False, run_on_click=True,
                )

if __name__ == "__main__":
    demo.queue(max_size=12).launch(mcp_server=True)