mage_v7.py 34 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753
  1. # ============================================================
  2. # Mage-VL 第七版
  3. # 实时洗衣机衣物分析
  4. #
  5. # 功能:
  6. # 1. 4 秒窗口实时分析
  7. # 2. StreamMind Gate 数值稳定修复
  8. # 3. BF16 优先 / T4 自动 FP16 + FP32 Gate fallback
  9. # 4. 累计衣物状态
  10. # 5. 动态洗涤程序
  11. # 6. 动态漂洗次数
  12. # 7. 视频结束后的最终综合洗涤方案
  13. #
  14. # 显存优化:
  15. # - Codec max_pixels = 48000
  16. # - target_canvas = 20
  17. # - max_new_tokens = 48
  18. # - 不保存 GPU 上的历史 inputs
  19. # - 每段结束主动清理 CUDA cache
  20. # ============================================================
  21. import os
  22. import gc
  23. import time
  24. import subprocess
  25. import importlib
  26. import warnings
  27. import json
  28. import math
  29. from pathlib import Path
  30. # ============================================================
  31. # 0. 环境
  32. # ============================================================
  33. os.environ.setdefault(
  34. "PYTORCH_CUDA_ALLOC_CONF",
  35. "expandable_segments:True,max_split_size_mb:128"
  36. )
  37. os.environ.setdefault(
  38. "ONLINE_CODEC_CACHE_DIR",
  39. "/tmp/mage_vl_codec_cache"
  40. )
  41. os.environ.setdefault(
  42. "STREAMMIND_CLIP_CACHE",
  43. "/tmp/mage_vl_clips"
  44. )
  45. os.makedirs(
  46. os.environ["ONLINE_CODEC_CACHE_DIR"],
  47. exist_ok=True
  48. )
  49. os.makedirs(
  50. os.environ["STREAMMIND_CLIP_CACHE"],
  51. exist_ok=True
  52. )
  53. # ============================================================
  54. # 1. Import
  55. # ============================================================
  56. import torch
  57. from transformers import (
  58. AutoModelForCausalLM,
  59. AutoProcessor,
  60. )
  61. # ============================================================
  62. # 2. 配置
  63. # ============================================================
  64. MODEL_ID = "microsoft/Mage-VL"
  65. VIDEO_FILE = "/content/test.mp4"
  66. # ---------- 实时窗口 ----------
  67. SEGMENT_SEC = 4.0
  68. MAX_SEGMENTS = None
  69. # 测试时可以改成:
  70. # MAX_SEGMENTS = 3
  71. # ---------- Gate ----------
  72. GATE_THRESHOLD = 0.50
  73. # ---------- LLM ----------
  74. MAX_NEW_TOKENS = 48
  75. # ---------- 最终总结 ----------
  76. FINAL_MAX_NEW_TOKENS = 192
  77. # ---------- Codec 显存 ----------
  78. # 第六版 64000
  79. # 第七版降低到 48000
  80. CODEC_MAX_PIXELS = 48_000
  81. # 第六版 24
  82. # 第七版降低到 20
  83. CODEC_TARGET_CANVAS = 20
  84. # ---------- 实时模拟 ----------
  85. SIMULATE_REALTIME = True
  86. # ============================================================
  87. # 3. CUDA
  88. # ============================================================
  89. if not torch.cuda.is_available():
  90. raise RuntimeError("未检测到 CUDA GPU")
  91. DEVICE = "cuda"
  92. GPU_NAME = torch.cuda.get_device_name(0)
  93. CUDA_VERSION = torch.version.cuda
  94. # ============================================================
  95. # 4. 自动选择 dtype
  96. # ============================================================
  97. BF16_SUPPORTED = False
  98. try:
  99. BF16_SUPPORTED = torch.cuda.is_bf16_supported()
  100. except Exception:
  101. BF16_SUPPORTED = False
  102. if BF16_SUPPORTED:
  103. MODEL_DTYPE = torch.bfloat16
  104. print("[dtype] GPU 支持 BF16")
  105. print("[dtype] Mage-VL = BF16")
  106. print("[dtype] StreamMind Gate = BF16")
  107. else:
  108. # T4 等旧 GPU:
  109. # 模型仍然使用 FP16 节省显存
  110. # Gate 单独尝试 FP32
  111. MODEL_DTYPE = torch.float16
  112. print("[dtype] GPU 不支持 BF16")
  113. print("[dtype] Mage-VL = FP16")
  114. print("[dtype] StreamMind Gate = FP32 fallback")
  115. print()
  116. print("=" * 60)
  117. print("Mage-VL 第七版")
  118. print("=" * 60)
  119. print(f"[GPU] {GPU_NAME}")
  120. print(f"[CUDA] {CUDA_VERSION}")
  121. print(f"[BF16] {BF16_SUPPORTED}")
  122. print(f"[MODEL DTYPE] {MODEL_DTYPE}")
  123. print(f"[CODEC] max_pixels={CODEC_MAX_PIXELS}")
  124. print(f"[CODEC] target_canvas={CODEC_TARGET_CANVAS}")
  125. print("=" * 60)
  126. # ============================================================
  127. # 5. Prompt
  128. # ============================================================
  129. PROMPT_TEMPLATE = """
  130. 你是洗衣机的实时视觉分析助手。
  131. 一个人正在不断把衣服放进洗衣机。
  132. 你的任务不是简单描述画面,而是:
  133. 【识别当前新增衣物】
  134. +
  135. 【维护累计衣物】
  136. +
  137. 【根据累计衣物动态选择最佳洗涤程序】。
  138. --------------------------------------------------
  139. 已有累计状态
  140. --------------------------------------------------
  141. {history}
  142. --------------------------------------------------
  143. 当前视频片段
  144. --------------------------------------------------
  145. 请只重点判断当前这个视频片段中新出现或正在加入洗衣机的衣物。
  146. 请分析:
  147. 1. 当前新增衣物
  148. - 类型:
  149. 牛仔裤、T恤、衬衫、毛衣、外套、
  150. 运动服、裙子、裤子、床品、大件衣物等
  151. - 颜色
  152. - 厚薄
  153. - 重量感
  154. - 可能材质:
  155. 棉、牛仔、涤纶、尼龙、羊毛、混纺等
  156. 如果无法确定材质,必须说“可能是”,不要假装确定。
  157. --------------------------------------------------
  158. 2. 累计衣物变化
  159. --------------------------------------------------
  160. 结合已有累计状态。
  161. 不要重复计算已经出现过的同一件衣服。
  162. 判断当前加入后,最终洗衣桶中的主要衣物组成。
  163. --------------------------------------------------
  164. 3. 动态洗涤程序
  165. --------------------------------------------------
  166. 根据目前累计衣物重新判断。
  167. 不要机械沿用上一段程序。
  168. 特别考虑:
  169. - 牛仔
  170. - 棉质
  171. - 化纤
  172. - 混纺
  173. - 毛衣
  174. - 针织
  175. - 外套
  176. - 厚重衣物
  177. - 大件衣物
  178. - 不同材质混洗
  179. - 深浅颜色混合
  180. 如果新增衣物导致程序应该改变,请明确指出。
  181. 例如:
  182. 牛仔 → 混合
  183. 混合 → 大物/混合
  184. 漂洗 2 次 → 3 次
  185. --------------------------------------------------
  186. 4. 当前推荐
  187. --------------------------------------------------
  188. 程序:
  189. 水温:
  190. 漂洗:
  191. 脱水:
  192. --------------------------------------------------
  193. 5. 判断原则
  194. --------------------------------------------------
  195. 最终程序优先保护最脆弱或最限制程序选择的衣物。
  196. 如果存在明显不适合混洗的衣物,需要提醒。
  197. 例如:
  198. 羊毛/针织 → 不建议与普通衣物混洗
  199. 深浅色混合 → 注意掉色
  200. 厚外套 + 普通衣物 → 优先考虑厚重衣物的洗涤要求
  201. --------------------------------------------------
  202. 输出要求
  203. --------------------------------------------------
  204. 简洁。
  205. 只输出:
  206. 【新增衣物】
  207. 【累计变化】
  208. 【当前程序】
  209. 【理由】
  210. 不要长篇解释。
  211. """.strip()
  212. # ============================================================
  213. # 6. 最终总结 Prompt
  214. # ============================================================
  215. FINAL_PROMPT = """
  216. 你是洗衣机最终洗涤决策模块。
  217. 下面是视频实时分析过程中得到的所有衣物状态:
  218. {history}
  219. 请综合整个视频。
  220. 不要只参考最后一个片段。
  221. --------------------------------------------------
  222. 任务
  223. --------------------------------------------------
  224. 1. 汇总最终实际放入洗衣机的衣物
  225. 2. 删除重复描述
  226. 3. 判断主要衣物类型
  227. 4. 判断主要可能材质
  228. 5. 找出最影响洗涤程序选择的衣物
  229. 6. 判断是否存在不适合混洗的衣物
  230. 7. 给出最终最佳洗涤程序
  231. --------------------------------------------------
  232. 洗涤程序选择原则
  233. --------------------------------------------------
  234. 如果只有牛仔类:
  235. 优先考虑牛仔程序。
  236. 如果普通棉质、T恤、衬衫等混合:
  237. 优先考虑混合程序。
  238. 如果出现较厚外套、大件或明显重量较大的衣物:
  239. 考虑大物/混合或更适合厚重衣物的程序。
  240. 如果存在羊毛、针织等特殊衣物:
  241. 优先保护特殊衣物,并明确提示不建议混洗。
  242. 漂洗次数:
  243. 普通衣物一般较少。
  244. 衣物较多、较厚、洗涤剂残留风险较高时增加漂洗。
  245. 脱水:
  246. 根据最终衣物厚薄和特殊材质判断。
  247. 不要为了保持前面的建议而保持原程序。
  248. 如果最终衣物组成发生变化,必须重新选择。
  249. --------------------------------------------------
  250. 输出格式
  251. --------------------------------------------------
  252. 【最终衣物】
  253. 列出主要衣物。
  254. 【主要材质】
  255. 列出主要可能材质。
  256. 【最终洗涤程序】
  257. 程序:
  258. 水温:
  259. 漂洗:
  260. 脱水:
  261. 【为什么】
  262. 用2~4句话说明最终选择。
  263. 【注意事项】
  264. 只列真正重要的问题。
  265. """.strip()
  266. # ============================================================
  267. # 7. 加载 Processor
  268. # ============================================================
  269. print()
  270. print("[init] 加载 Processor...")
  271. processor = AutoProcessor.from_pretrained(
  272. MODEL_ID,
  273. trust_remote_code=True,
  274. )
  275. # ============================================================
  276. # 8. 加载 Mage-VL
  277. # ============================================================
  278. print("[init] 加载 Mage-VL...")
  279. model = AutoModelForCausalLM.from_pretrained(
  280. MODEL_ID,
  281. trust_remote_code=True,
  282. torch_dtype=MODEL_DTYPE,
  283. attn_implementation="sdpa",
  284. low_cpu_mem_usage=True,
  285. ).to(DEVICE).eval()
  286. print("[init] 模型加载完成")
  287. # ============================================================
  288. # 9. Codec
  289. # ============================================================
  290. _REMOTE_PKG = type(processor).__module__.rsplit(".", 1)[0]
  291. codec_mod = importlib.import_module(
  292. _REMOTE_PKG + ".codec_video_processing_mage_vl"
  293. )
  294. CodecConfig = codec_mod.CodecConfig
  295. def codec_config():
  296. override = {
  297. "engine": "hevc",
  298. "target_canvas": CODEC_TARGET_CANVAS,
  299. "patch": 16,
  300. }
  301. kwargs = dict(
  302. processor._codec_config_defaults
  303. )
  304. dcvc = dict(
  305. kwargs.get("dcvc") or {}
  306. )
  307. kwargs.update(override)
  308. if dcvc:
  309. kwargs["dcvc"] = dcvc
  310. kwargs["max_pixels"] = CODEC_MAX_PIXELS
  311. return CodecConfig(**kwargs), override
  312. # ============================================================
  313. # 10. 加载 Gate
  314. # ============================================================
  315. print("[gate] 加载 Cognition Gate...")
  316. try:
  317. gate = model.model._load_streammind_gate()
  318. print("[gate] Cognition Gate 加载成功")
  319. try:
  320. gate_dtype = next(
  321. gate.parameters()
  322. ).dtype
  323. print(
  324. f"[gate] 当前 dtype = {gate_dtype}"
  325. )
  326. except Exception:
  327. pass
  328. except Exception as e:
  329. gate = None
  330. print(
  331. f"[gate] 加载失败: {type(e).__name__}: {e}"
  332. )
  333. # ============================================================
  334. # 11. T4 / FP16 下:
  335. # 尝试把 Gate 单独提升到 FP32
  336. # ============================================================
  337. if (
  338. gate is not None
  339. and not BF16_SUPPORTED
  340. ):
  341. try:
  342. gate = gate.float()
  343. model.model.streammind_gate = gate
  344. print(
  345. "[gate] FP32 fallback 已启用"
  346. )
  347. except Exception as e:
  348. print(
  349. "[gate] FP32 fallback 失败:",
  350. e
  351. )
  352. # ============================================================
  353. # 12. 工具
  354. # ============================================================
  355. def duration(path):
  356. try:
  357. r = subprocess.run(
  358. [
  359. "ffprobe",
  360. "-v",
  361. "error",
  362. "-show_entries",
  363. "format=duration",
  364. "-of",
  365. "default=nw=1:nk=1",
  366. str(path),
  367. ],
  368. capture_output=True,
  369. text=True,
  370. check=True,
  371. timeout=60,
  372. )
  373. return float(
  374. r.stdout.strip()
  375. )
  376. except Exception:
  377. return 0.0
  378. # ============================================================
  379. # 13. Prompt
  380. # ============================================================
  381. def make_prompt(history_text=""):
  382. text = PROMPT_TEMPLATE.format(
  383. history=(
  384. history_text
  385. if history_text
  386. else "暂无,这是第一个视频片段。"
  387. )
  388. )
  389. messages = [
  390. {
  391. "role": "user",
  392. "content": [
  393. {
  394. "type": "video"
  395. },
  396. {
  397. "type": "text",
  398. "text": text,
  399. },
  400. ],
  401. }
  402. ]
  403. return processor.apply_chat_template(
  404. messages,
  405. tokenize=False,
  406. add_generation_prompt=True,
  407. )
  408. # ============================================================
  409. # 14. 视频切片
  410. # ============================================================
  411. def extract_clip(
  412. src,
  413. start,
  414. length,
  415. out,
  416. ):
  417. tmp = out.with_name(
  418. ".tmp_" + out.name
  419. )
  420. if tmp.exists():
  421. try:
  422. tmp.unlink()
  423. except Exception:
  424. pass
  425. subprocess.run(
  426. [
  427. "ffmpeg",
  428. "-y",
  429. "-loglevel",
  430. "error",
  431. "-ss",
  432. f"{start:.3f}",
  433. "-i",
  434. str(src),
  435. "-t",
  436. f"{length:.3f}",
  437. "-c:v",
  438. "libx264",
  439. "-preset",
  440. "ultrafast",
  441. "-crf",
  442. "26",
  443. "-an",
  444. "-pix_fmt",
  445. "yuv420p",
  446. str(tmp),
  447. ],
  448. check=True,
  449. capture_output=True,
  450. timeout=180,
  451. )
  452. tmp.replace(out)
  453. return out
  454. # ============================================================
  455. # 15. CPU → GPU
  456. # ============================================================
  457. def to_cuda(inputs):
  458. result = {}
  459. for k, v in inputs.items():
  460. if not hasattr(v, "to"):
  461. continue
  462. v = v.to(DEVICE)
  463. if k == "pixel_values":
  464. v = v.to(MODEL_DTYPE)
  465. result[k] = v
  466. return result
  467. # ============================================================
  468. # 16. 清理 GPU
  469. # ============================================================
  470. def cleanup(*objs):
  471. for x in objs:
  472. try:
  473. del x
  474. except Exception:
  475. pass
  476. gc.collect()
  477. try:
  478. torch.cuda.empty_cache()
  479. torch.cuda.ipc_collect()
  480. except Exception:
  481. pass
  482. # ============================================================
  483. # 17. GPU 显存
  484. # ============================================================
  485. def gpu_mem():
  486. allocated = (
  487. torch.cuda.memory_allocated()
  488. / 1024**3
  489. )
  490. reserved = (
  491. torch.cuda.memory_reserved()
  492. / 1024**3
  493. )
  494. peak = (
  495. torch.cuda.max_memory_allocated()
  496. / 1024**3
  497. )
  498. return (
  499. f"GPU={allocated:.2f}GB "
  500. f"/ reserved={reserved:.2f}GB "
  501. f"/ peak={peak:.2f}GB"
  502. )
  503. # ============================================================
  504. # 18. Gate
  505. # ============================================================
  506. @torch.inference_mode()
  507. def run_gate(inputs):
  508. feat = {}
  509. for k in (
  510. "pixel_values",
  511. "image_grid_thw",
  512. "patch_positions",
  513. ):
  514. if k not in inputs:
  515. continue
  516. v = inputs[k].to(DEVICE)
  517. if k == "pixel_values":
  518. v = v.to(MODEL_DTYPE)
  519. feat[k] = v
  520. logits = None
  521. try:
  522. # ----------------------------------------------------
  523. # BF16 GPU
  524. # ----------------------------------------------------
  525. if BF16_SUPPORTED:
  526. logits = (
  527. model.model
  528. .streammind_gate_forward_segments(
  529. [feat]
  530. )
  531. )
  532. # ----------------------------------------------------
  533. # T4 / FP16 fallback
  534. #
  535. # 不直接使用官方 gate forward,
  536. # 而是:
  537. #
  538. # visual token → FP32 → Gate FP32
  539. #
  540. # 尽量避免 FP16 Gate overflow
  541. # ----------------------------------------------------
  542. else:
  543. visual_tokens = (
  544. model.model
  545. ._streammind_vision_tokens(
  546. feat["pixel_values"],
  547. feat["image_grid_thw"],
  548. patch_positions=feat.get(
  549. "patch_positions"
  550. ),
  551. )
  552. )
  553. if not torch.isfinite(
  554. visual_tokens
  555. ).all():
  556. print(
  557. "[gate] visual tokens "
  558. "出现 NaN/Inf"
  559. )
  560. return 0.0
  561. visual_tokens = (
  562. visual_tokens.float()
  563. )
  564. logits = gate(
  565. torch.cat(
  566. [visual_tokens],
  567. dim=1,
  568. )
  569. )
  570. # ----------------------------------------------------
  571. # 官方 Gate 输出:
  572. #
  573. # [B,T,2]
  574. #
  575. # 0 = silence
  576. # 1 = response
  577. # ----------------------------------------------------
  578. if not isinstance(
  579. logits,
  580. torch.Tensor
  581. ):
  582. raise RuntimeError(
  583. "Gate 输出不是 Tensor"
  584. )
  585. if logits.ndim != 3:
  586. raise RuntimeError(
  587. "Unexpected gate logits shape: "
  588. + str(tuple(logits.shape))
  589. )
  590. gate_logits = logits[0]
  591. # Codec 时间长度
  592. length = int(
  593. inputs[
  594. "image_grid_thw"
  595. ][:, 0]
  596. .sum()
  597. .item()
  598. )
  599. if length <= 0:
  600. raise RuntimeError(
  601. f"Invalid gate length={length}"
  602. )
  603. pos = min(
  604. length - 1,
  605. gate_logits.shape[0] - 1,
  606. )
  607. last_logits = (
  608. gate_logits[pos]
  609. .float()
  610. )
  611. # ----------------------------------------------------
  612. # 关键:
  613. # 在 softmax 之前检查
  614. # ----------------------------------------------------
  615. if not torch.isfinite(
  616. last_logits
  617. ).all():
  618. print(
  619. "[gate] WARNING:"
  620. " logits =",
  621. last_logits
  622. .detach()
  623. .cpu()
  624. .tolist(),
  625. )
  626. return 0.0
  627. # ----------------------------------------------------
  628. # 用 log_softmax 更稳定
  629. # ----------------------------------------------------
  630. log_prob = torch.log_softmax(
  631. last_logits,
  632. dim=-1,
  633. )
  634. prob = torch.exp(
  635. log_prob[1]
  636. )
  637. if not torch.isfinite(prob):
  638. print(
  639. "[gate] WARNING:"
  640. " probability NaN/Inf"
  641. )
  642. return 0.0
  643. return float(
  644. prob.item()
  645. )
  646. except Exception as e:
  647. print(
  648. "[gate] ERROR:",
  649. type(e).__name__,
  650. e,
  651. )
  652. return 0.0
  653. finally:
  654. cleanup(
  655. feat,
  656. logits,
  657. )
  658. # ============================================================
  659. # 19. Codec + Gate
  660. # ============================================================
  661. @torch.inference_mode()
  662. def analyze_clip(
  663. clip,
  664. prompt,
  665. codec_override,
  666. ):
  667. with warnings.catch_warnings():
  668. warnings.simplefilter(
  669. "ignore"
  670. )
  671. inputs = processor(
  672. text=[prompt],
  673. videos=[
  674. str(clip)
  675. ],
  676. video_backend="codec",
  677. max_pixels=CODEC_MAX_PIXELS,
  678. codec_config=codec_override,
  679. return_tensors="pt",
  680. padding=False,
  681. )
  682. prob = run_gate(
  683. inputs
  684. )
  685. return (
  686. inputs,
  687. prob,
  688. )
  689. # ============================================================
  690. # 20. LLM
  691. # ============================================================
  692. @torch.inference_mode()
  693. def generate(
  694. inputs,
  695. max_new_tokens=MAX_NEW_TOKENS,
  696. ):
  697. out = model.generate(
  698. **inputs,
  699. max_new_tokens=max_new_tokens,
  700. do_sample=False,
  701. use_cache=True,
  702. )
  703. n = (
  704. inputs[
  705. "input_ids"
  706. ].shape[1]
  707. )
  708. text = processor.tokenizer.decode(
  709. out[0, n:],
  710. skip_special_tokens=True,
  711. ).strip()
  712. del out
  713. return text
  714. # ============================================================
  715. # 21. 最终总结
  716. # ============================================================
  717. @torch.inference_mode()
  718. def final_summary(history):
  719. if not history:
  720. return (
  721. "没有获得有效的衣物分析结果。"
  722. )
  723. history_text = "\n\n".join(
  724. history
  725. )
  726. prompt = FINAL_PROMPT.format(
  727. history=history_text
  728. )
  729. inputs = processor.tokenizer(
  730. prompt,
  731. return_tensors="pt",
  732. truncation=True,
  733. max_length=4096,
  734. )
  735. inputs = {
  736. k: v.to(DEVICE)
  737. for k, v in inputs.items()
  738. }
  739. try:
  740. out = model.generate(
  741. **inputs,
  742. max_new_tokens=FINAL_MAX_NEW_TOKENS,
  743. do_sample=False,
  744. use_cache=True,
  745. )
  746. n = (
  747. inputs[
  748. "input_ids"
  749. ].shape[1]
  750. )
  751. result = (
  752. processor.tokenizer.decode(
  753. out[0, n:],
  754. skip_special_tokens=True,
  755. )
  756. .strip()
  757. )
  758. return result
  759. finally:
  760. cleanup(
  761. inputs
  762. )
  763. try:
  764. del out
  765. except Exception:
  766. pass
  767. # ============================================================
  768. # 22. 实时视频分析
  769. # ============================================================
  770. def stream_video_analysis(
  771. video_path,
  772. segment_sec=SEGMENT_SEC,
  773. gate_threshold=GATE_THRESHOLD,
  774. max_segments=MAX_SEGMENTS,
  775. max_new_tokens=MAX_NEW_TOKENS,
  776. simulate_realtime=SIMULATE_REALTIME,
  777. ):
  778. path = Path(
  779. video_path
  780. )
  781. if not path.exists():
  782. raise FileNotFoundError(
  783. video_path
  784. )
  785. total = duration(
  786. path
  787. )
  788. if total <= 0:
  789. raise RuntimeError(
  790. "无法获取视频时长"
  791. )
  792. clip_dir = Path(
  793. os.environ[
  794. "STREAMMIND_CLIP_CACHE"
  795. ]
  796. )
  797. clip_dir.mkdir(
  798. parents=True,
  799. exist_ok=True,
  800. )
  801. # --------------------------------------------------------
  802. # 两个历史:
  803. #
  804. # results:
  805. # 给用户看的完整输出
  806. #
  807. # history:
  808. # 给模型看的累计状态
  809. # --------------------------------------------------------
  810. results = []
  811. history = []
  812. start = 0.0
  813. index = 0
  814. yield (
  815. f"开始实时分析 | "
  816. f"时长={total:.1f}s | "
  817. f"窗口={segment_sec:.1f}s | "
  818. f"Gate={gate_threshold:.2f} | "
  819. f"Codec={CODEC_MAX_PIXELS}px"
  820. )
  821. # --------------------------------------------------------
  822. # Codec
  823. # --------------------------------------------------------
  824. _, codec_override = codec_config()
  825. # ========================================================
  826. # 主循环
  827. # ========================================================
  828. while start < total - 0.05:
  829. if (
  830. max_segments is not None
  831. and index >= max_segments
  832. ):
  833. break
  834. end = min(
  835. total,
  836. start + segment_sec,
  837. )
  838. length = (
  839. end - start
  840. )
  841. index += 1
  842. t0 = time.time()
  843. clip = clip_dir / (
  844. f"{path.stem}_"
  845. f"{index:06d}_"
  846. f"{start:.2f}_"
  847. f"{end:.2f}.mp4"
  848. )
  849. inputs = None
  850. cuda_inputs = None
  851. said = None
  852. try:
  853. # ------------------------------------------------
  854. # 1. 当前累计历史
  855. # ------------------------------------------------
  856. history_text = "\n".join(
  857. history[-6:]
  858. )
  859. prompt = make_prompt(
  860. history_text
  861. )
  862. # ------------------------------------------------
  863. # 2. 精确切片
  864. # ------------------------------------------------
  865. extract_clip(
  866. path,
  867. start,
  868. length,
  869. clip,
  870. )
  871. # ------------------------------------------------
  872. # 3. Codec + Gate
  873. # ------------------------------------------------
  874. inputs, prob = analyze_clip(
  875. clip,
  876. prompt,
  877. codec_override,
  878. )
  879. head = (
  880. f"**{start:.1f}–"
  881. f"{end:.1f}s** "
  882. f"Gate=`{prob:.2f}`"
  883. )
  884. # ------------------------------------------------
  885. # 4. Gate
  886. #
  887. # 注意:
  888. #
  889. # 本项目不让 Gate 决定“是否识别衣服”。
  890. #
  891. # Gate 只是告诉我们:
  892. # 这个片段是否具有明显的事件变化。
  893. #
  894. # 为避免漏掉新衣物:
  895. # 每个片段仍然进入 Mage-VL。
  896. # ------------------------------------------------
  897. if prob < gate_threshold:
  898. elapsed = time.time() - t0
  899. speed = length / max(elapsed, 0.001)
  900. line = (
  901. f"{head} → 🔇 跳过"
  902. f"(Gate < {gate_threshold:.2f})"
  903. f" | Total={elapsed:.2f}s"
  904. )
  905. results.append(line)
  906. yield "\n\n".join(results)
  907. # --------------------------------------------
  908. # 关键:
  909. #
  910. # 不进入 Mage-VL
  911. # 不 generate
  912. # 不加入 history
  913. #
  914. # 直接进入下一个 4 秒片段
  915. # --------------------------------------------
  916. continue
  917. # ------------------------------------------------
  918. # Gate 通过
  919. # ------------------------------------------------
  920. yield (
  921. "\n\n".join(results)
  922. + (
  923. "\n\n"
  924. if results
  925. else ""
  926. )
  927. + f"{head} → 🧠 分析衣物中..."
  928. )
  929. gt = time.time()
  930. # ------------------------------------------------
  931. # GPU
  932. # ------------------------------------------------
  933. cuda_inputs = to_cuda(
  934. inputs
  935. )
  936. # ------------------------------------------------
  937. # Mage-VL
  938. # ------------------------------------------------
  939. said = generate(
  940. cuda_inputs,
  941. max_new_tokens=max_new_tokens,
  942. )
  943. gen_time = (
  944. time.time()
  945. - gt
  946. )
  947. elapsed = (
  948. time.time()
  949. - t0
  950. )
  951. speed = (
  952. length
  953. / max(elapsed, 0.001)
  954. )
  955. # ------------------------------------------------
  956. # 保存结果
  957. # ------------------------------------------------
  958. line = (
  959. f"{head} → 🗣️ "
  960. f"{said}\n"
  961. f" Generate="
  962. f"{gen_time:.2f}s | "
  963. f"Total="
  964. f"{elapsed:.2f}s | "
  965. f"Realtime="
  966. f"{speed:.2f}×"
  967. )
  968. results.append(
  969. line
  970. )
  971. history.append(
  972. f"{start:.1f}-"
  973. f"{end:.1f}s:"
  974. f"{said}"
  975. )
  976. yield "\n\n".join(
  977. results
  978. )
  979. # if prob < gate_threshold:
  980. # gate_state = (
  981. # "Gate较低"
  982. # )
  983. # else:
  984. # gate_state = (
  985. # "Gate触发"
  986. # )
  987. # # ------------------------------------------------
  988. # # 5. 进入 Mage-VL
  989. # # ------------------------------------------------
  990. # yield (
  991. # "\n\n".join(
  992. # results
  993. # )
  994. # + (
  995. # "\n\n"
  996. # if results
  997. # else ""
  998. # )
  999. # + f"{head} "
  1000. # + f"→ 🧠 {gate_state},"
  1001. # + "分析衣物中..."
  1002. # )
  1003. # gt = time.time()
  1004. # # ------------------------------------------------
  1005. # # GPU
  1006. # # ------------------------------------------------
  1007. # cuda_inputs = to_cuda(
  1008. # inputs
  1009. # )
  1010. # # ------------------------------------------------
  1011. # # Generate
  1012. # # ------------------------------------------------
  1013. # said = generate(
  1014. # cuda_inputs,
  1015. # max_new_tokens=max_new_tokens,
  1016. # )
  1017. # gen_time = (
  1018. # time.time()
  1019. # - gt
  1020. # )
  1021. # elapsed = (
  1022. # time.time()
  1023. # - t0
  1024. # )
  1025. # speed = (
  1026. # length
  1027. # / max(elapsed, 0.001)
  1028. # )
  1029. # # ------------------------------------------------
  1030. # # 6. 当前结果
  1031. # # ------------------------------------------------
  1032. # line = (
  1033. # f"{head} → 🗣️ "
  1034. # f"{said}\n"
  1035. # f" Generate="
  1036. # f"{gen_time:.2f}s | "
  1037. # f"Total="
  1038. # f"{elapsed:.2f}s | "
  1039. # f"Realtime="
  1040. # f"{speed:.2f}×"
  1041. # )
  1042. # results.append(
  1043. # line
  1044. # )
  1045. # # ------------------------------------------------
  1046. # # 7. 只把模型结果加入累计历史
  1047. # # ------------------------------------------------
  1048. # history.append(
  1049. # f"{start:.1f}-"
  1050. # f"{end:.1f}s:"
  1051. # f"{said}"
  1052. # )
  1053. # # ------------------------------------------------
  1054. # # 8. 实时输出
  1055. # # ------------------------------------------------
  1056. # yield "\n\n".join(
  1057. # results
  1058. # )
  1059. # ------------------------------------------------
  1060. # 9. 模拟真实播放速度
  1061. # ------------------------------------------------
  1062. elapsed = (
  1063. time.time()
  1064. - t0
  1065. )
  1066. if (
  1067. simulate_realtime
  1068. and elapsed < length
  1069. ):
  1070. time.sleep(
  1071. length - elapsed
  1072. )
  1073. except Exception as e:
  1074. line = (
  1075. f"**{start:.1f}–"
  1076. f"{end:.1f}s** → "
  1077. f"❌ "
  1078. f"{type(e).__name__}: "
  1079. f"{e}"
  1080. )
  1081. results.append(
  1082. line
  1083. )
  1084. yield "\n\n".join(
  1085. results
  1086. )
  1087. finally:
  1088. # ------------------------------------------------
  1089. # 清理 GPU
  1090. # ------------------------------------------------
  1091. cleanup(
  1092. cuda_inputs,
  1093. inputs,
  1094. said,
  1095. )
  1096. # ------------------------------------------------
  1097. # 删除临时视频
  1098. # ------------------------------------------------
  1099. if clip.exists():
  1100. try:
  1101. clip.unlink()
  1102. except Exception:
  1103. pass
  1104. start = end
  1105. # ------------------------------------------------
  1106. # 显存统计
  1107. # ------------------------------------------------
  1108. print(
  1109. f"[segment {index}] "
  1110. f"{start:.1f}/"
  1111. f"{total:.1f}s | "
  1112. f"{gpu_mem()}",
  1113. flush=True,
  1114. )
  1115. # ------------------------------------------------
  1116. # 重置 peak
  1117. # ------------------------------------------------
  1118. try:
  1119. torch.cuda.reset_peak_memory_stats()
  1120. except Exception:
  1121. pass
  1122. # ========================================================
  1123. # 视频结束
  1124. # ========================================================
  1125. yield (
  1126. "\n\n".join(
  1127. results
  1128. )
  1129. + "\n\n"
  1130. + "========================================\n"
  1131. + "🔄 视频播放完成,正在综合全部衣物...\n"
  1132. + "========================================"
  1133. )
  1134. # ========================================================
  1135. # 最终综合
  1136. # ========================================================
  1137. try:
  1138. final = final_summary(
  1139. history
  1140. )
  1141. final_output = (
  1142. "\n\n".join(
  1143. results
  1144. )
  1145. + "\n\n"
  1146. + "========================================\n"
  1147. + "🧺 最终综合洗涤方案\n"
  1148. + "========================================\n\n"
  1149. + final
  1150. + "\n\n"
  1151. + f"✅ 完成,共 {index} 个片段"
  1152. )
  1153. yield final_output
  1154. except Exception as e:
  1155. yield (
  1156. "\n\n".join(
  1157. results
  1158. )
  1159. + "\n\n"
  1160. + "========================================\n"
  1161. + "🧺 最终综合洗涤方案\n"
  1162. + "========================================\n\n"
  1163. + f"❌ 综合分析失败:"
  1164. + f"{type(e).__name__}: {e}"
  1165. )
  1166. # ============================================================
  1167. # 23. Main
  1168. # ============================================================
  1169. if __name__ == "__main__":
  1170. for output in stream_video_analysis(
  1171. VIDEO_FILE,
  1172. segment_sec=SEGMENT_SEC,
  1173. gate_threshold=GATE_THRESHOLD,
  1174. max_segments=MAX_SEGMENTS,
  1175. max_new_tokens=MAX_NEW_TOKENS,
  1176. simulate_realtime=SIMULATE_REALTIME,
  1177. ):
  1178. # 清屏
  1179. print(
  1180. "\033[H\033[J",
  1181. end=""
  1182. )
  1183. print(
  1184. "=== Mage-VL 第七版:"
  1185. "洗衣机实时衣物分析 ===\n"
  1186. )
  1187. print(output)
  1188. print(
  1189. "\n"
  1190. + "-" * 60
  1191. )