app.py

import os, json, time, subprocess, tempfile, shutil
import gradio as gr
import spaces
import torch
from functools import lru_cache
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

# --- 配置區 ---
REPO_ID = "Memories-ai/security_model"
TOKEN = os.environ.get("HF_TOKEN")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "512"))  # 原 768 太高，先收斂
FORCE_FPS = int(os.environ.get("FORCE_FPS", "1"))              # 影片抽幀 6fps 足夠 caption
TARGET_MAX_W = int(os.environ.get("TARGET_MAX_W", "480"))     # 寬度上限 1280 (<=720p)
DEBUG_TIMINGS = os.environ.get("DEBUG_TIMINGS", "0") == "1"    # 1 時把分段時間附在輸出

# 速度小優化（Ampere 以後有效）
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# --- Debug 模式控制 ---
DEBUG = os.environ.get("DEBUG", "0") == "1"

def dprint(*args, **kwargs):
    """只在 DEBUG 模式下印出訊息"""
    if DEBUG:
        print(*args, **kwargs)

# ---------- 實用工具：ffprobe & 可能轉碼 ----------
def _run_quiet(cmd: list[str]):
    subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

def ffprobe_meta(path: str):
    try:
        out = subprocess.check_output([
            "ffprobe","-v","error","-select_streams","v:0",
            "-show_entries","stream=codec_name,width,height,avg_frame_rate",
            "-of","json", path
        ])
        data = json.loads(out.decode("utf-8"))
        st = data["streams"][0] if data.get("streams") else {}
        fps = 0.0
        afr = st.get("avg_frame_rate","0/0")
        if isinstance(afr,str) and "/" in afr:
            num, den = afr.split("/")
            fps = float(num)/float(den) if float(den) != 0 else 0.0
        return {
            "codec": st.get("codec_name"),
            "w": int(st.get("width") or 0),
            "h": int(st.get("height") or 0),
            "fps": fps
        }
    except Exception:
        return {"codec": None, "w": 0, "h": 0, "fps": 0.0}

def maybe_transcode(input_path: str):
    """
    碰到 HEVC/H.265 或解析度太大時，快速轉成 H.264 + yuv420p + 目標寬度 + 限制 FPS
    轉完回傳 (path, used_temp=True/False, reason)
    """
    meta = ffprobe_meta(input_path)
    codec, w, h, fps = meta["codec"], meta["w"], meta["h"], meta["fps"]

    need_codec_fix = codec in ("hevc","h265")
    need_resize = (w and w > TARGET_MAX_W)
    need_fps = (fps and fps > FORCE_FPS + 0.5)

    if not (need_codec_fix or need_resize or need_fps):
        return input_path, False, {"meta": meta, "transcoded": False}

    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    out_path = tmp.name; tmp.close()

    # scale 只在寬度超標時啟動，保留比例；fps 超標則限速
    vf_parts = []
    if need_resize:
        vf_parts.append(f"scale='min({TARGET_MAX_W},iw)':-2")
    if need_fps:
        vf_parts.append(f"fps={FORCE_FPS}")
    vf = ",".join(vf_parts) if vf_parts else "scale=trunc(iw/2)*2:trunc(ih/2)*2"

    cmd = [
        "ffmpeg","-y","-i", input_path,
        "-vsync","vfr",
        "-c:v","libx264","-preset","veryfast","-crf","23",
        "-pix_fmt","yuv420p",
        "-vf", vf,
        "-c:a","aac","-b:a","128k",
        "-movflags","+faststart",
        out_path
    ]
    _run_quiet(cmd)
    return out_path, True, {"meta": meta, "transcoded": True, "vf": vf}

# ---------- 分段計時 ----------
class Timer:
    def __init__(self): self.t0=time.perf_counter(); self.spans=[]
    def mark(self, name, dur): self.spans.append((name, round(dur,3)))
    def result(self):
        total = round(time.perf_counter()-self.t0, 3)
        return {"total_s": total, **{k:v for k,v in self.spans}}

# 載入模型（用私有 token），自動上 GPU
@lru_cache(maxsize=1)
def _load_model_and_processor():
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        REPO_ID,
        torch_dtype="auto",
        device_map="auto",
        token=TOKEN,
    )
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
    return model, processor

# 你的提示詞（可直接沿用）
PROBLEM_TEXT = """You are an expert AI agent specializing in smart home security video analysis.

Your task is to **generate a single, detailed description (caption)** for a video captured by a household security camera with a perfect field of view. Your description should identify and document any potentially risky, suspicious, or anomalous situations as defined by smart home surveillance criteria.

---

### **Description Requirements**

### 1. **Core Event Analysis & Classification**

Your primary task is to identify the main event in the video. Your description should be centered around this event, providing a clear and neutral account of all visible activities. Pay special attention to and be prepared to identify events from the following categories:

- **Violent & Criminal Activity**: Abuse, Arrest, Arson, Assault, Burglary, Explosion, Fighting, Robbery, Shooting, Shoplifting, Stealing, Vandalism
- **Safety & Environmental Risks**: Road Accidents, Wildlife Intrusions
- **Other Anomalies**: Unusual Pet Behavior, or other abnormal events.
- **Normal Activity**: If no anomalies are present, classify the footage as a Normal Video.

### 2. **Human Subject Descriptions**

- **Count & Basic Info**: Count and describe **each person** in the footage.
- **Individual Details**: For each individual, provide:
    - **Gender** (if visually inferable)
    - **Age group** (e.g., child, adult, elderly)
    - **Clothing details**: hat, glasses, top, pants.
- **Identity Concealment**: Clearly note any attempts to hide identity (e.g., wearing masks, hoodies, hats).
- **Specific Actions & Behaviors**:
    - **General Behavior**: Describe suspicious actions like loitering, trespassing, or fleeing the scene.
    - **Aggressive Actions**: Detail physical attacks, shoving, striking, the use of weapons (e.g., firearms, knives), or any form of **Fighting** or **Assault**.
    - **Coercion & Force**: Report any acts of forcibly taking property from a person (**Robbery**) or interactions involving physical coercion (**Abuse**).
    - **Law Enforcement Actions**: Only label individuals as law enforcement if wearing **clearly identifiable uniforms and caps**. Describe their actions, such as using restraints (e.g., handcuffs), making an **Arrest**, and escorting individuals.

### 3. **Interaction with Environment**

- **Routine Interactions**: Document interactions with **doors, packages, bins, fences, pets, or furniture**.
- **Property Crime & Damage**:
    - **Unlawful Entry & Theft**: Highlight any lock-picking, window breaking, attempted or successful entry into a structure (**Burglary**), and the act of taking items from the premises (**Stealing**).
    - **Vandalism**: Describe any acts of intentional property damage, such as graffiti, breaking windows, or damaging structures.
    - **Arson & Explosions**: Detail the act of starting a fire, using accelerants, and any visual signs of an **Explosion** (e.g., a bright flash, smoke, debris).

### 4. **Animal & Vehicle Descriptions**

- **Animals**: Describe wild or domestic animals, including **species, appearance, behavior**, and their interactions with **people or property**. Report any risks or damage caused (e.g., tipping trash bins).
- **Vehicles**: If a vehicle is present, provide its **color, type, brand (if visible), and license plate (if readable)**. Describe its relation to the scene (e.g., arriving, a person exiting, involvement in a collision (**Road Accident**)).

### 5. **Reporting Style & Format**

- **Spatio-Temporal Consistency**:
    - Describe events in **chronological order**.
    - Use **landmarks and camera orientation** to describe locations (e.g., “enters from the left side of the frame,” “approaches the trash bin near the driveway”).
    - Include **time context** if inferable (e.g., “at night,” “in daylight”).
- **Objective Surveillance Tone**:
    - Use a **neutral, concise, and factual tone**.
    - **Avoid speculation or storytelling**. Do not infer intent or emotions beyond visible behavior.

---

### **Example Output (Based on New Requirements)**

In daylight, an unmarked white van enters from the right side of the frame and parks at the curb. An adult male exits the driver's seat. The man is wearing a black hooded sweatshirt with the hood up, obscuring his face, dark pants, and white sneakers. He walks directly to the house, pulls a tool from his pocket, and begins prying at the front door lock. After approximately 15 seconds, the door opens, and the man enters the residence. Three minutes later, he exits the house carrying a laptop and a backpack. He quickly returns to the van and drives away, exiting the frame to the left. No other people or vehicles are visible. This event can be categorized as a burglary.
"""

# [NEW] 自訂 prompt 的長度（避免爆 context）
MAX_PROMPT_CHARS = int(os.environ.get("MAX_PROMPT_CHARS", "8000"))

# [NEW] 清理/截斷使用者自訂 prompt
def _sanitize_prompt(p: str) -> str:
    if not p:
        return ""
    p = p.strip()
    if len(p) > MAX_PROMPT_CHARS:
        p = p[:MAX_PROMPT_CHARS]
    return p

@spaces.GPU(duration=900)
# [CHANGED] 增加可選參數 prompt_text；預設空字串代表客戶未提供
def caption_video(video_path: str, prompt_text: str = "") -> str:
    if not video_path:
        return "No video provided."

    T = Timer()
    
    t = time.perf_counter()
    model, processor = _load_model_and_processor()
    if torch.cuda.is_available(): torch.cuda.synchronize()
    T.mark("load_model_s", time.perf_counter()-t)

    print("[ENV] MAX_NEW_TOKENS =", MAX_NEW_TOKENS, flush=True)
    dprint("[CUDA] available =", torch.cuda.is_available(), flush=True)
    if torch.cuda.is_available():
        dprint("[CUDA] device =", torch.cuda.get_device_name(0), flush=True)
    try:
        dprint("[MODEL] device_map =", getattr(model, "hf_device_map", None), flush=True)
        dprint("[MODEL] first_param_device =", next(model.parameters()).device, flush=True)
    except Exception as e:
        dprint("[MODEL] device inspect error:", e, flush=True)    
    
    
    # 1) 可能轉碼 / 降維 / 限 FPS
    t = time.perf_counter()
    safe_path, used_temp, tr_info = maybe_transcode(video_path)
    T.mark("maybe_transcode_s", time.perf_counter()-t)

    # [NEW] 若客戶沒輸入，就用後端的 PROBLEM_TEXT；否則使用自訂（但不顯示預設給客戶）
    user_prompt = _sanitize_prompt(prompt_text)
    selected_prompt = user_prompt if user_prompt else PROBLEM_TEXT    
    
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "video", "video": video_path, "fps": 1},
                {"type": "text", "text": selected_prompt},  # [CHANGED] 由 selected_prompt 餵入
            ],
        }
    ]

    # 建構聊天模板與多模態輸入
    t = time.perf_counter()    
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs, video_kwargs = process_vision_info(
        messages, return_video_kwargs=True
    )
    
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
        **video_kwargs,
    )

    ctx_len = int(inputs.input_ids.shape[-1])
    try:
        n_frames = len(video_inputs[0]) if isinstance(video_inputs, (list, tuple)) else None
    except Exception:
        n_frames = None
    dprint(f"[CTX] tokens={ctx_len}, est_frames={n_frames}", flush=True)

    # 上 GPU（若可）
    if torch.cuda.is_available():
        inputs = inputs.to("cuda")
        torch.cuda.synchronize()
    T.mark("preprocess_s", time.perf_counter()-t)
        

    gen_kwargs = dict(
        max_new_tokens=MAX_NEW_TOKENS,
        do_sample=False,           # caption 任務較適合確定性解碼，速度更快
        temperature=0.0,
        top_p=1.0,
        use_cache=True,
    )
    
    dprint("[ENV] MAX_NEW_TOKENS =", MAX_NEW_TOKENS, flush=True)
    
    t = time.perf_counter()
    with torch.inference_mode():
        generated_ids = model.generate(**inputs, **gen_kwargs)
        if torch.cuda.is_available(): torch.cuda.synchronize()
    T.mark("generate_s", time.perf_counter()-t)

    gen_time = time.perf_counter() - t
    dprint(f"[GEN] time_s={gen_time:.3f}", flush=True)    


    # 5) 後處理
    t = time.perf_counter()
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )
    if torch.cuda.is_available(): torch.cuda.synchronize()
    T.mark("postprocess_s", time.perf_counter()-t)

    # print
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    n_new = int(generated_ids_trimmed[0].shape[-1]) if generated_ids_trimmed else 0
    tps = (n_new / gen_time) if gen_time > 0 else 0.0
    dprint(f"[GEN] new_tokens={n_new}, tps={tps:.3f} tok/s", flush=True)
    

    # 6) 清理暫存檔
    if used_temp:
        try: os.remove(safe_path)
        except Exception: pass    

    # 打印詳細 timing 到日誌（HF Spaces Logs 可見）
    rt = T.result()
    known = sum(v for k, v in rt.items() if k != "total_s")
    other = round(rt["total_s"] - known, 3)
    dprint({"timings": rt | {"other_s": other}, "transcode": tr_info}, flush=True)

    caption = (output_text[0] if output_text else "").strip()
    if DEBUG_TIMINGS:
        caption += (
            f"\n\n[timings] total={rt['total_s']}s, "
            f"transcode={rt.get('maybe_transcode_s','-')}s, "
            f"preprocess={rt.get('preprocess_s','-')}s, "
            f"generate={rt.get('generate_s','-')}s, "
            f"other={other}s"
        )
    return caption

# Gradio 介面
# demo = gr.Interface(
#     fn=caption_video,
#     inputs=gr.Video(label="Upload a video (mp4, mov, etc.)"),
#     outputs=gr.Textbox(label="Caption", lines=18),
#     title="Smart Home Video Caption (Private weights)",
#     description="Upload a short clip to generate a factual surveillance-style caption.",
#     allow_flagging="never",
# )

# [CHANGED] Gradio 介面：不再把預設 prompt 顯示給客戶
# - 第二個輸入改成「可選的自訂 prompt」文字框，value=""，僅 placeholder 提示
# - 也可放在 Accordion 內，避免一般使用者注意到
with gr.Blocks(title="Smart Home Video Caption (Private weights)") as demo:
    gr.Markdown("Upload a short clip to generate a factual surveillance-style caption.")
    with gr.Row():
        video_in = gr.Video(label="Upload a video (mp4, mov, etc.)")
    with gr.Accordion("Advanced (optional custom prompt)", open=False):  # [NEW]
        prompt_in = gr.Textbox(
            label="Custom prompt (optional)",
            value="",                 # 空字串 => 客戶看不到預設內容
            placeholder="Leave empty to use the default internal prompt",
            lines=10,
            show_copy_button=True
        )
    caption_out = gr.Textbox(label="Caption", lines=18)

    run_btn = gr.Button("Run")
    run_btn.click(fn=caption_video, inputs=[video_in, prompt_in], outputs=[caption_out])


if __name__ == "__main__":
    demo.launch(show_error=True)