"""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)