app.py

"""Peak-End-Net: Video Aesthetic Assessment Demo.

Upload a video and get an overall aesthetic score plus ten fine-grained
attribute scores (composition, shotsize, lighting, visualtone, color,
depthoffield, expression, movement, costume, makeup), based on the
peak-end rule from cognitive psychology.

Model: https://huggingface.co/GD-ML/Peak-End-Net
Paper:  https://arxiv.org/abs/2607.13941
Code:   https://github.com/AMAP-ML/Peak-End-Net
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any CUDA-touching import
import torch
import numpy as np
import gradio as gr

# --- Patch torch.load to allow loading the legacy checkpoint (weights_only) ---
_orig_torch_load = torch.load
def _patched_load(*a, **k):
    k.setdefault("weights_only", False)
    return _orig_torch_load(*a, **k)
torch.load = _patched_load

# --- Import the Peak-End-Net model architecture ---
from models.gated_fusion import GatedFusionPeakAesNetV4
from modules.rawvideo_util import RawVideoExtractor

# --- Download the checkpoint from the Hub ---
from huggingface_hub import hf_hub_download

MODEL_REPO = "GD-ML/Peak-End-Net"
CHECKPOINT_FILENAME = "Peak-End-Net.pth"

# 11 aesthetic attribute names (index 0 = overall).
SCORE_NAMES = [
    "overall", "composition", "shotsize", "lighting",
    "visualtone", "color", "depthoffield",
    "expression", "movement", "costume", "makeup",
]

MAX_FRAMES = 12
IMAGE_RESOLUTION = 224


def load_model():
    """Build the model skeleton and load the merged self-contained weights."""
    print(f"[app] Downloading checkpoint from {MODEL_REPO}...")
    ckpt_path = hf_hub_download(MODEL_REPO, CHECKPOINT_FILENAME)

    print(f"[app] Building model (downloads CLIP ViT-L/14 weights from OpenAI)...")
    model = GatedFusionPeakAesNetV4(
        stage1_checkpoint_path=None,
        ava_checkpoint_path=None,
        pretrained_clip_name="ViT-L/14",
        max_frames=MAX_FRAMES,
        frame_context_window=3,
        end_window=3,
        rhythm_dim=64,
        end_ratio=0.25,
        scoring_hidden_dim=256,
        fusion_hidden_dim=128,
    )

    print(f"[app] Loading checkpoint from {ckpt_path}...")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    missing, unexpected = model.load_state_dict(state_dict, strict=True)
    if missing:
        print(f"[app] [warn] missing keys: {len(missing)}")
    if unexpected:
        print(f"[app] [warn] unexpected keys: {len(unexpected)}")

    model.to("cuda").eval()
    print("[app] Model loaded and moved to CUDA.")
    return model


# Load at module scope (ZeroGPU intercepts .to("cuda"))
model = load_model()


def preprocess_video(video_path, max_frames=MAX_FRAMES, image_resolution=IMAGE_RESOLUTION, slice_framepos=2):
    """Decode a video into the model input tensor (same pipeline as training).

    Returns:
        video:      [1, 1, T, 1, 3, H, W] float tensor
        video_mask: [1, T] long tensor
    """
    extractor = RawVideoExtractor(framerate=1, size=image_resolution)

    video = np.zeros(
        (1, max_frames, 1, 3, image_resolution, image_resolution),
        dtype=np.float32,
    )
    video_mask = np.zeros((1, max_frames), dtype=np.int64)
    max_len = 0

    raw = extractor.get_video_data(video_path)["video"]
    if len(raw.shape) > 3:
        raw_slice = extractor.process_raw_data(raw)
        if max_frames < raw_slice.shape[0]:
            if slice_framepos == 0:
                sel = raw_slice[:max_frames, ...]
            elif slice_framepos == 1:
                sel = raw_slice[-max_frames:, ...]
            else:
                idx = np.linspace(0, raw_slice.shape[0] - 1, num=max_frames, dtype=int)
                sel = raw_slice[idx, ...]
        else:
            sel = raw_slice
        sel = extractor.process_frame_order(sel, frame_order=0)
        max_len = sel.shape[0]
        if max_len >= 1:
            video[0][:max_len, ...] = sel

    video_mask[0][:max_len] = 1

    # [1, T, 1, 3, H, W] -> [1, 1, T, 1, 3, H, W] (batch dim)
    video_tensor = torch.tensor(video).unsqueeze(0)
    mask_tensor = torch.tensor(video_mask)
    return video_tensor, mask_tensor


@spaces.GPU(duration=60)
def assess_video(video_path: str) -> dict:
    """Assess the aesthetic quality of a video.

    Args:
        video_path: Path to a video file (mp4, avi, mov, etc.).

    Returns:
        A dictionary with the overall aesthetic score, ten fine-grained
        attribute scores, and the fusion gate details.
    """
    if video_path is None:
        return {"error": "Please upload a video file."}

    video, video_mask = preprocess_video(video_path)
    video = video.to("cuda")
    video_mask = video_mask.to("cuda")

    with torch.no_grad():
        final_scores, aux = model(video, video_mask)
    scores = final_scores[0].cpu().numpy()

    result = {}
    for name, value in zip(SCORE_NAMES, scores):
        result[name] = round(float(value), 4)

    result["gate"] = round(float(aux["gate"][0].item()), 4)
    result["S_model"] = round(float(aux["S_model"][0].item()), 4)
    result["S_static"] = round(float(aux["S_static"][0].item()), 4)

    return result


# --- Gradio UI ---

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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # 🎬 Peak-End-Net: Video Aesthetic Assessment

        Upload a video to get an **overall aesthetic score** and **ten fine-grained attribute scores**
        based on the *peak-end rule* from cognitive psychology — the idea that humans judge
        experiences by their most intense moments and their ending rather than a simple average.

        **Model:** [GD-ML/Peak-End-Net](https://huggingface.co/GD-ML/Peak-End-Net) ·
        **Paper:** [arXiv:2607.13941](https://arxiv.org/abs/2607.13941) ·
        **Code:** [GitHub](https://github.com/AMAP-ML/Peak-End-Net)
        """)

        with gr.Row():
            video_input = gr.Video(label="Upload Video", sources=["upload"])
            run_btn = gr.Button("Assess Aesthetic", variant="primary")

        output_json = gr.JSON(label="Aesthetic Scores")

        with gr.Accordion("Score Details", open=False):
            gr.Markdown("""
            **Score meanings** (all on a ~0–10 scale):

            | Score | Description |
            |---|---|
            | **overall** | Overall aesthetic quality (gated fusion of model + static frame scores) |
            | **composition** | Quality of visual composition and framing |
            | **shotsize** | Appropriateness of shot size / framing distance |
            | **lighting** | Quality and mood of lighting |
            | **visualtone** | Overall visual tone and atmosphere |
            | **color** | Color palette and grading quality |
            | **depthoffield** | Depth of field and focus quality |
            | **expression** | Emotional expression and impact (human subjects) |
            | **movement** | Quality of camera and subject movement |
            | **costume** | Costume and styling quality (human subjects) |
            | **makeup** | Makeup quality (human subjects) |

            **Fusion details:**
            - **gate**: The fusion gate weight (0 = trust static frame average, 1 = trust the model)
            - **S_model**: Stage-1 model overall score
            - **S_static**: Average per-frame AVA aesthetic score
            - Final overall = gate × S_model + (1 − gate) × S_static
            """)

        gr.Examples(
            examples=[
                ["examples/landscape.mp4"],
                ["examples/ocean_waves.mp4"],
            ],
            inputs=[video_input],
            outputs=output_json,
            fn=assess_video,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=assess_video,
        inputs=[video_input],
        outputs=output_json,
        api_name="assess",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)