app.py

import gradio as gr
import torch
import io
from PIL import Image
import numpy as np
import spaces  # Import spaces for ZeroGPU compatibility
import math
import re
from einops import rearrange
from mmengine.config import Config
from xtuner.registry import BUILDER 

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from scripts.camera.cam_dataset import Cam_Generator
from scripts.camera.visualization.visualize_batch import make_perspective_figures

from huggingface_hub import snapshot_download
import os
local_path = snapshot_download(
    repo_id="KangLiao/Puffin",
    repo_type="model",       
    #filename="Puffin-Base.pth",         
    local_dir="checkpoints/",  
    local_dir_use_symlinks=False,     
    revision="main",                  
)

local_path_vae = snapshot_download(
    repo_id="wusize/Puffin",
    repo_type="model",       
    #filename="Puffin-Base.pth",         
    local_dir="checkpoints_vae/",  
    local_dir_use_symlinks=False,     
    revision="main",                  
)


NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)

def center_crop(image):
    w, h = image.size
    s = min(w, h)
    l = (w - s) // 2
    t = (h - s) // 2
    return image.crop((l, t, l + s, t + s))


##### load model
# base model
config = "configs/pipelines/stage_2_base.py"
config = Config.fromfile(config)
model = BUILDER.build(config.model).cuda().bfloat16().eval()
checkpoint_path = "checkpoints/Puffin-Base.pth"
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint, strict=False)

checkpoint_path_vae = "checkpoints_vae/vae.pth"
checkpoint_vae = torch.load(checkpoint_path_vae)
model.vae.load_state_dict(checkpoint_vae, strict=False)


# thinking model
config_thinking = "configs/pipelines/stage_3_thinking.py"
config_thinking = Config.fromfile(config_thinking)
model_think = BUILDER.build(config_thinking.model).cuda().bfloat16().eval()
checkpoint_path = "checkpoints/Puffin-Thinking.pth"
checkpoint = torch.load(checkpoint_path)
model_think.load_state_dict(checkpoint, strict=False)
model_think.vae.load_state_dict(checkpoint_vae, strict=False)

description = r"""
<b>Official Gradio demo</b> for <a href='https://kangliao929.github.io/projects/puffin/' target='_blank'><b>Thinking with Camera: A Unified Multimodal Model for Camera-Centric Understanding and Generation</b></a>.<br>
🔥 We make the first attempt to integrate camera geometry into a unified multimodal model, introducing a camera-centric framework (<b>Puffin</b>) to advance multimodal spatial intelligence.<br>
🖼️ Try to switch the task table and choose different prompts or images to get the generation or understanding results.<br>
"""

article = r"""<h3>
<b>If Puffin is helpful, please help to star the <a href='https://github.com/KangLiao929/Puffin' target='_blank'>Github Repo</a>. Thank you.</b></h3>
<hr>

📑 **Citation**
<br>
If our work is useful for your research, please consider citing:
```bibtex
@article{liao2025puffin,
title={Thinking with Camera: A Unified Multimodal Model for Camera-Centric Understanding and Generation},
author={Liao, Kang and Wu, Size and Wu, Zhonghua and Jin, Linyi and Wang, Chao and Wang, Yikai and Wang, Fei and Li, Wei and Loy, Chen Change},
journal={arXiv preprint arXiv:2510.08673},
year={2025}
}
```
📧 **Contact**
<br>
If you have any questions, please feel free to reach me out at <b>kang.liao@ntu.edu.sg</b>.
<br>
"""


import base64
with open("assets/Puffin.png", "rb") as f:
    img_bytes = f.read()
img_b64 = base64.b64encode(img_bytes).decode()

html_img = f'''
<div style="display:flex; justify-content:center; align-items:center; width:100%;">
    <img src="data:image/png;base64,{img_b64}" style="border:none; width:150px; height:auto;"/>
</div>
'''

@torch.inference_mode()
@spaces.GPU(duration=120) 
# Multimodal Understanding function
def camera_understanding(image_src, thinking_und, question, seed, progress=gr.Progress(track_tqdm=True)):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    # torch.manual_seed(seed)
    # np.random.seed(seed)
    # torch.cuda.manual_seed(seed)
    print(torch.cuda.is_available())

    prompt = ("Describe the image in detail. Then reason its spatial distribution and estimate its camera parameters (roll, pitch, and field-of-view).")
    if thinking_und:
        prompt = ("Reason the spatial distribution of this image in a thinking mode, and then estimate its camera parameters (roll, pitch, and field-of-view).")

    image = Image.fromarray(image_src).convert('RGB')
    image = center_crop(image)
    image = image.resize((512, 512))
    x = torch.from_numpy(np.array(image)).float()
    x = x / 255.0
    x = 2 * x - 1
    x = rearrange(x, 'h w c -> c h w')

    with torch.no_grad():
        if thinking_und:
            outputs = model_think.understand(prompt=[prompt], pixel_values=[x], progress_bar=False)
        else:
            outputs = model.understand(prompt=[prompt], pixel_values=[x], progress_bar=False)

    text = outputs[0]
    gen = Cam_Generator(mode="cot") if thinking_und else Cam_Generator(mode="base")
    cam = gen.get_cam(text)
    
    bgr = np.array(image)[:, :, ::-1].astype(np.float32) / 255.0
    rgb = bgr[:, :, ::-1].copy()        
    image_tensor = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
    single_batch = {
        "image": image_tensor,
        "up_field": cam[:2].unsqueeze(0),
        "latitude_field": cam[2:].unsqueeze(0),
    }

    figs = make_perspective_figures(single_batch, single_batch, n_pairs=1)
    saved_paths = []
    save_dir = "temp/"
    os.makedirs(save_dir, exist_ok=True)
    
    for k, fig in figs.items():
        if "up_field" in k:
            suffix = "_up"
        elif "latitude_field" in k:
            suffix = "_lat"
        else:
            suffix = f"_{k}"
        out_path = os.path.join(save_dir, f"camera_map_vis{suffix}.png")
        plt.tight_layout()
        fig.savefig(out_path, dpi=200, bbox_inches='tight', pad_inches=0)
        plt.close(fig)
        saved_paths.append(out_path)

    img_up = Image.open(saved_paths[0]).convert("RGB")
    img_lat = Image.open(saved_paths[1]).convert("RGB")
    w, h = img_up.size
    left = max(0, w - h)
    img_up = img_up.crop((left, 0, w, h))
    w, h = img_lat.size
    left = max(0, w - h)
    img_lat = img_lat.crop((left, 0, w, h))
    img_up = img_up.resize((512, 512))
    img_lat = img_lat.resize((512, 512))
    
    gap = 10
    W, H = img_up.size
    combined = Image.new("RGB", (W * 2 + gap, H), (255, 255, 255))
    combined.paste(img_up, (0, 0))
    combined.paste(img_lat, (W + gap, 0))

    return text, combined


@torch.inference_mode()
@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
def generate_image(prompt_scene,
                   seed=42,
                   roll=0.1,
                   pitch=0.1,
                   fov=1.0,
                   thinking_gen=False,
                   progress=gr.Progress(track_tqdm=True)):
    # Clear CUDA cache and avoid tracking gradients
    torch.cuda.empty_cache()
    # Set the seed for reproducible results
    # if seed is not None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    print(torch.cuda.is_available())
    
    generator = torch.Generator().manual_seed(seed)
    prompt_camera = (
        "The camera parameters (roll, pitch, and field-of-view) are: "
        f"{roll:.4f}, {pitch:.4f}, {fov:.4f}."
    )
    
    prompt_thinking = ("Given a scene description and corresponding camera parameters, "
                       "merge them into a coherent prompt and generate an accurate visualization "
                       "that highlights visual cues for spatial reasoning.")
    gen = Cam_Generator()
    cam_map = gen.get_cam(prompt_camera).to(model.device)
    cam_map = cam_map / (math.pi / 2)
    
    prompt = prompt_scene + " " + prompt_camera
    
    bsz = 4
    with torch.no_grad():
        if thinking_gen:
            images, output_reasoning = model_think.generate(
                prompt=[prompt]*bsz,
                cfg_prompt=[""]*bsz,
                pixel_values_init=None,
                cfg_scale=4.5,
                num_steps=50,
                cam_values=[[cam_map]]*bsz,
                progress_bar=False,
                reasoning=thinking_gen,
                prompt_reasoning=[prompt_thinking]*bsz,
                generator=generator,
                height=512,
                width=512
            )
        else:
            images, output_reasoning = model.generate(
                prompt=[prompt]*bsz,
                cfg_prompt=[""]*bsz,
                pixel_values_init=None,
                cfg_scale=4.5,
                num_steps=50,
                cam_values=[[cam_map]]*bsz,
                progress_bar=False,
                reasoning=thinking_gen,
                prompt_reasoning=[prompt_thinking]*bsz,
                generator=generator,
                height=512,
                width=512
            )

        images = rearrange(images, 'b c h w -> b h w c')
        images = torch.clamp(127.5 * images + 128.0, 0, 255).to("cpu", dtype=torch.uint8).numpy()
        ret_images = [Image.fromarray(image) for image in images]
        return ret_images, output_reasoning[0]


# Gradio interface
css = '''
.gradio-container {max-width: 960px !important}
'''

custom_css = """
#input-image {
    aspect-ratio: 1 / 1;
    width: 100%;
    max-width: 100%;
    height: auto;
    display: flex;
    align-items: center;
    justify-content: center;
}
#input-image img {
    max-width: 100%;
    max-height: 100%;
    object-fit: contain;
    display: block;
}
#main-columns {
    gap: 60px; 
}
#main-columns > .gr-column {
    flex: 1; 
}
#compare-image {
    width: 100%;
    aspect-ratio: 1 / 1; 
    display: flex;
    align-items: center;
    justify-content: center;
    margin: 0;
    padding: 0;
    max-width: 100%;
    box-sizing: border-box;
}
#compare-image svg.svelte-zyxd38 {
    position: absolute !important; 
    top: 50% !important;           
    left: 50% !important;          
    transform: translate(-50%, -50%) !important; 
}
#compare-image .icon.svelte-1oiin9d {
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
}
#compare-image {
    position: relative;
    overflow: hidden;
}
.new_button {background-color: #171717 !important; color: #ffffff !important; border: none !important;}
.new_button:hover {background-color: #4b4b4b !important;}
#start-button {
    background: linear-gradient(135deg, #2575fc 0%, #6a11cb 100%);
    color: white;
    border: none;
    padding: 12px 24px;
    font-size: 16px;
    font-weight: bold;
    border-radius: 12px;
    cursor: pointer;
    box-shadow: 0 0 12px rgba(100, 100, 255, 0.7);
    transition: all 0.3s ease;
}
#start-button:hover {
    transform: scale(1.05);
    box-shadow: 0 0 20px rgba(100, 100, 255, 1);
}
<style>
.button-wrapper {
    width: 30%;
    text-align: center; 
}
.wide-button {
    width: 83% !important;
    background-color: black !important;
    color: white !important;
    border: none !important;
    padding: 8px 0 !important;
    font-size: 16px !important;
    display: inline-block;
    margin: 30px 0px 0px 50px ;
}
.wide-button:hover {
    background-color: #656262 !important;
}
</style>
"""

with gr.Blocks(css=custom_css) as demo:
    #gr.Markdown("# Puffin")
    gr.HTML(html_img)
    gr.Markdown(description)

    with gr.Tab("Camera-controllable Generation"):
        gr.Markdown(value="## Camera-controllable Generation")

        prompt_input = gr.Textbox(label="Scene prompt")

        with gr.Accordion("Camera parameters (in radius)", open=True):
            with gr.Row():
                roll = gr.Slider(minimum=-0.7854, maximum=0.7854, value=0.1000, step=0.1000, label="roll value")
                pitch = gr.Slider(minimum=-0.7854, maximum=0.7854, value=-0.1000, step=0.1000, label="pitch value")
                fov = gr.Slider(minimum=0.3491, maximum=1.8326, value=1.5000, step=0.1000, label="fov value")
        with gr.Accordion("Settings", open=True):
            with gr.Row(equal_height=True):
                thinking_gen = gr.Radio(
                    ["Thinking"],
                    label=None,
                    value=None,
                    show_label=False,
                    interactive=True
                )

                seed_input = gr.Number(
                    label="Seed (Optional)",
                    precision=0,
                    value=42
                )
            
        generation_button = gr.Button("Generate Images")
    
        image_output = gr.Gallery(label="Generated images", columns=4, rows=1)
        
        output_reasoning = gr.Textbox(label="Response (only in thinking)")
    
        examples_t2i = gr.Examples(
            label="Prompt examples",
            examples=[
                "A sunny day casts light on two warmly colored buildings—yellow with green accents and deeper orange—framed by a lush green tree, with a blue sign and street lamp adding details in the foreground.",
                "A high-vantage-point view of lush, autumn-colored mountains blanketed in green and gold, set against a clear blue sky with scattered white clouds, offering a tranquil and breathtaking vista of a serene valley below.",
                "A grand, historic castle with pointed spires and elaborate stone structures stands against a clear blue sky, flanked by a circular fountain, vibrant red flowers, and neatly trimmed hedges in a beautifully landscaped garden.",
                "A serene aerial view of a coastal landscape at sunrise/sunset, featuring warm pink and orange skies transitioning to cool blues, with calm waters stretching to rugged, snow-capped mountains in the background, creating a tranquil and picturesque scene.",
                "A worn, light-yellow walls room with herringbone terracotta floors and three large arched windows framed in pink trim and white panes, showcasing signs of age and disrepair, overlooks a residential area through glimpses of greenery and neighboring buildings.",
                "The Milky Way rises above a remote bay near the grassy shores of Iceland, its silvery arc mirrored in the still, glassy waters below."
            ],
            inputs=prompt_input,
        )

    with gr.Tab("Camera Understanding"):
        gr.Markdown(value="## Camera Understanding")
        image_input = gr.Image()
        
        with gr.Accordion("Settings", open=True):            
            with gr.Row(equal_height=True):
                thinking_und = gr.Radio(
                    ["Thinking"],
                    label=None,
                    value=None,
                    show_label=False,
                    interactive=True
                )

                und_seed_input = gr.Number(
                    label="Seed (Optional)",
                    precision=0,
                    value=42
                )

        understanding_button = gr.Button("Chat")
        understanding_output = gr.Textbox(label="Response")
        
        camera_map = gr.Image(label="Camera maps (up vector & latitude)")

        examples_inpainting = gr.Examples(
            label="Examples",
            examples=[
                "assets/1.jpg",
                "assets/2.jpg",
                "assets/3.jpg",
                "assets/4.jpg",
                "assets/5.jpg",
                "assets/6.jpg",
            ],
            inputs=image_input,
        )

    generation_button.click(
        fn=generate_image,
        inputs=[prompt_input, seed_input, roll, pitch, fov, thinking_gen],
        outputs=[image_output, output_reasoning]
    )

    understanding_button.click(
        camera_understanding,
        inputs=[image_input, thinking_und, und_seed_input],
        outputs=[understanding_output, camera_map]
    )
    
    gr.Markdown(article)

demo.launch(share=True)