Update app.py

app.py

@@ -12,63 +12,50 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ------------------------------------------------------------------------
-"""URSA TI2V application —— 纯 CPU + 4bit (bitsandbytes NF4) 版本.
-
-相对 ZeroGPU 版的改动：
-  1. 去掉 CUDA / ZeroGPU 依赖，全部在 CPU 上跑（spaces.GPU 变 no-op）。
-  2. 权重量化：transformer 里的 nn.Linear 换成 bitsandbytes Linear4bit (NF4)。
-  3. CPU 上没有 flex_attention，把 FlexAttentionCausal2D 改写成 SDPA 实现。
-  4. 默认分辨率 / 帧数 / 步数降到 CPU 能承受的档位。
-"""
+"""URSA TI2V application (torchao int4 / int8 weight-only quantization)."""
 
 # =========================
-# ✅ 第 0 步：环境变量（必须早于 import torch，线程数只能在这里设）
+# ✅ 第 1 步：环境变量必须最先设置（早于 torch / spaces / diffnext）
 # =========================
 import os
 
-_THREADS = str(os.cpu_count() or 4)
-os.environ.setdefault("OMP_NUM_THREADS", _THREADS)
-os.environ.setdefault("MKL_NUM_THREADS", _THREADS)
-# 纯 CPU 模式：想用 GPU 就自行 export CUDA_VISIBLE_DEVICES=0
-os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
-# torch.compile 在 CPU 上要走 inductor，URSA 只用它编译 RoPE / flex_attention，收益为负
-os.environ["TORCH_COMPILE_DISABLE"] = "1"
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
+# 只在 token 存在时设置，避免把 HF_TOKEN 污染成空字符串
 if os.environ.get("HF_TOKEN"):
     os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", os.environ["HF_TOKEN"])
 os.environ["HF_HUB_ENABLE_HUGGINGFACE_TRANSFER"] = "1"
 os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+# 必须在 CUDA 初始化 / torch 导入前设置
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+# 关掉 TorchDynamo（torch 在 import 时读取该环境变量）
+os.environ["TORCH_COMPILE_DISABLE"] = "1"
 
-import argparse  # noqa: E402
-import shutil  # noqa: E402
-import time  # noqa: E402
-
-import numpy as np  # noqa: E402
-import PIL.Image  # noqa: E402
-import torch  # noqa: E402
-import torch.nn.functional as F  # noqa: E402
-
-torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", _THREADS)))
-try:
-    CPU_CAP = torch.backends.cpu.get_cpu_capability()
-except Exception:  # 老版本 torch 没有这个 API
-    CPU_CAP = "UNKNOWN"
-
-
-def cpu_supports_bf16() -> bool:
-    """CPU 是否能比较体面地跑 bf16（AVX512-BF16 / AMX）。"""
-    try:  # bitsandbytes 的判断最准，优先用
-        from bitsandbytes.functional import has_avx512bf16
+# =========================
+# ✅ 第 2 步：import spaces 必须早于 torch（ZeroGPU 硬性要求）
+# =========================
+import spaces
 
-        return bool(has_avx512bf16())
-    except Exception:
-        return CPU_CAP == "AVX512"
+import argparse
+import re
+import sys
+import time
 
+import gradio as gr
+import numpy as np
+import PIL.Image
+import torch
 
 # =========================
-# ✅ 第 1 步：把 torch.compile 变成恒等变换
-#    diffnext 在【模块导入时】就用 @torch.compile 装饰了 RoPE 实现，
-#    CPU 上 inductor 只会更慢（且首次编译会卡几十秒）。
+# 🔧 第 3 步：把 torch.compile 变成恒等变换
+#    diffnext 在【模块导入时】就用 @torch.compile 装饰了 RoPE 实现：
+#      - FlexRotaryEmbedding.PEFunc.interleaved_impl / partitioned_impl
+#      - RotaryEmbed3D.PEFunc.call_impl
+#      - FlexAttn.get_attn_func() 运行期调用 torch.compile(flex_attention)
+#    该编译是惰性的，直到首次真实推理（进入 @spaces.GPU、真卡挂载后）才真正触发，
+#    此时 TorchInductor 用 CUDA event 给 Triton kernel 做 autotune 微基准，
+#    在 ZeroGPU 的 fork worker 里 event 计时失效 →
+#    RuntimeError: Both events must be completed before calculating elapsed time
+#    ZeroGPU 官方明确不支持 torch.compile（每任务新起进程，JIT 无法复用）。
 # =========================
 try:
     torch._dynamo.config.disable = True
@@ -78,7 +65,9 @@ except Exception:
 
 def _eager_compile(model=None, **kwargs):
     """@torch.compile(...) / torch.compile(model) 都直接返回原对象。"""
-    return model if model is not None else (lambda fn: fn)
+    if model is not None:
+        return model
+    return lambda fn: fn
 
 
 torch.compile = _eager_compile
@@ -88,7 +77,8 @@ from diffnext.utils import export_to_image, export_to_video  # noqa: E402
 
 
 # =========================
-# 🔧 第 2 步：兜底 —— 把被编译的 RoPE 方法换成纯 eager 实现
+# 🔧 第 4 步：兜底 —— 直接把被编译的 RoPE 方法换成纯 eager 实现
+#    （即使第 3 步失效，这一步也能保证不会走 Inductor）
 # =========================
 def _patch_rope_eager():
     from diffnext.models import embeddings as emb
@@ -115,16 +105,23 @@ def _patch_rope_eager():
     return patched
 
 
+try:
+    _ROPE_PATCHED = _patch_rope_eager()
+except Exception as e:  # 不因打补丁失败而阻塞启动
+    _ROPE_PATCHED = ["FAILED: %r" % (e,)]
+
+print("[patch] torch.compile -> identity:", torch.compile is _eager_compile, flush=True)
+print("[patch] rope eager:", _ROPE_PATCHED, flush=True)
+
+
 # =========================
-# 🔧 第 3 步：修复 Qwen3Config 缺失 rope_theta
+# 🔧 第 5 步：修复 Qwen3Config 缺失 rope_theta
 # =========================
 def _patch_diffnext_qwen3():
     try:
         from diffnext.models.text_encoders import qwen3
     except Exception:
         return
-    if not hasattr(qwen3, "Qwen3RotaryEmbedding"):
-        return
 
     orig_from_config = qwen3.Qwen3RotaryEmbedding.from_config
 
@@ -132,7 +129,10 @@ def _patch_diffnext_qwen3():
     def patched_from_config(cls, config):
         if not hasattr(config, "rope_theta"):
             rp = getattr(config, "rope_parameters", None)
-            config.rope_theta = rp["rope_theta"] if isinstance(rp, dict) and "rope_theta" in rp else 1e6
+            if isinstance(rp, dict) and "rope_theta" in rp:
+                config.rope_theta = rp["rope_theta"]
+            else:
+                config.rope_theta = float(1e6)
         if not hasattr(config, "max_position_embeddings"):
             config.max_position_embeddings = int(getattr(config, "max_seq_len", 2048))
         return orig_from_config(config)
@@ -140,362 +140,356 @@ def _patch_diffnext_qwen3():
     qwen3.Qwen3RotaryEmbedding.from_config = patched_from_config
 
 
-# =========================
-# 🔧 第 4 步（CPU 关键）：flex_attention 在 CPU 上不可用，换成 SDPA
-#    Qwen3SdpaAttention 里是这么分支的：
-#        if self.flex_attn and self.flex_attn.offsets: -> flex_attention
-#        else:                                          -> F.scaled_dot_product_attention
-#    而推理 pipeline 从不调用 set_offsets()，所以正常就走 SDPA 分支。
-#    这里仍然重写 forward，保证万一 offsets 被设置也不会掉进 CPU 不支持的
-#    flex_attention（此时用一张物化出来的 bool mask 表达 2D 因果图）。
-# =========================
-_HAS_GQA = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2]) >= (2, 5)
-
-
-def _sdpa(q, k, v, attn_mask=None):
-    if _HAS_GQA:
-        return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=True)
-    n_rep = q.shape[1] // k.shape[1]  # 老 torch：手动展开 GQA
-    k, v = k.repeat_interleave(n_rep, dim=1), v.repeat_interleave(n_rep, dim=1)
-    return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
-
-
-def _patch_flex_attn_sdpa():
-    from diffnext.models.flex_attention import FlexAttentionCausal2D
-
-    def _build_mask(self, q):
-        """把 2D 因果规则物化成 [1, 1, T, T] 的 bool mask（只在 offsets 非空时用）。"""
-        offsets = torch.as_tensor(self.offsets, device=q.device)
-        flags = list(self.flags or [])
-        key = (tuple(self.offsets), tuple(flags), q.shape[-2], str(q.device))
-        if getattr(self, "_mask_key", None) == key:
-            return self._mask_cache
-        counts = offsets[1:] - offsets[:-1]
-        ids = torch.arange(len(counts), device=q.device).repeat_interleave(counts)
-        flags = torch.as_tensor(flags + [-1] * (len(counts) - len(flags)), device=q.device)
-        flags = flags.repeat_interleave(counts)
-        idx = torch.arange(ids.numel(), device=q.device)
-        # (qi >= ki) | (ids[qi] * flags[qi] == ids[ki])
-        mask = (idx[:, None] >= idx[None, :]) | ((ids[:, None] * flags[:, None]) == ids[None, :])
-        self._mask_key, self._mask_cache = key, mask[None, None]
-        return self._mask_cache
-
-    def forward(self, q, k, v):
-        if not self.offsets:  # 推理默认路径：全双向 SDPA
-            return _sdpa(q, k, v)
-        return _sdpa(q, k, v, attn_mask=_build_mask(self, q))
-
-    FlexAttentionCausal2D.forward = forward
-    return True
+_patch_diffnext_qwen3()
 
 
-try:
-    _ROPE_PATCHED = _patch_rope_eager()
-except Exception as e:
-    _ROPE_PATCHED = ["FAILED: %r" % (e,)]
-_patch_diffnext_qwen3()
-try:
-    _FLEX_PATCHED = _patch_flex_attn_sdpa()
-except Exception as e:
-    _FLEX_PATCHED = "FAILED: %r" % (e,)
+def parse_args():
+    parser = argparse.ArgumentParser(description="Serve URSA TI2V application")
+    parser.add_argument("--model", default="BAAI/URSA-0.6B-FSQ320", help="model path")
+    parser.add_argument("--device", type=int, default=0, help="device index")
+    parser.add_argument("--precision", default="float16", help="compute precision")
+    # ---- torchao 量化开关 ----
+    parser.add_argument(
+        "--quant",
+        default="int4",
+        choices=["none", "int4", "int8"],
+        help="torchao weight-only 量化：int4=tinygemm 4bit / int8=A16W8",
+    )
+    parser.add_argument(
+        "--quant-group-size",
+        type=int,
+        default=128,
+        choices=[32, 64, 128, 256],
+        help="分组量化粒度，越小越精细、压缩率越低（32 精度最好但量化最慢）",
+    )
+    parser.add_argument(
+        "--quant-algo",
+        default="hqq",
+        choices=["hqq", "tinygemm"],
+        help="int4 的 qparams 算法：hqq 精度高但量化慢；tinygemm 快但需 mslk 依赖",
+    )
+    parser.add_argument(
+        "--quant-keep-layers",
+        type=int,
+        default=2,
+        help="首尾各 N 层保持原精度（0=全部量化），越大越稳、越费显存",
+    )
+    parser.add_argument("--quant-lm-head", action="store_true", help="连 codebook 输出头一起量化（不建议）")
+    parser.add_argument("--quant-vae", action="store_true", help="连 VAE 一起量化（强烈不建议，画质损失大）")
+    return parser.parse_args()
 
-print(f"[patch] torch.compile -> identity : {torch.compile is _eager_compile}", flush=True)
-print(f"[patch] rope eager                 : {_ROPE_PATCHED}", flush=True)
-print(f"[patch] flex_attention -> sdpa     : {_FLEX_PATCHED}", flush=True)
-print(f"[env  ] cpu={CPU_CAP} threads={torch.get_num_threads()} bf16_fast={cpu_supports_bf16()}", flush=True)
+
+def crop_image(image, target_h, target_w):
+    h, w = image.height, image.width
+    aspect_ratio_target, aspect_ratio = target_w / target_h, w / h
+    if aspect_ratio > aspect_ratio_target:
+        new_w = int(h * aspect_ratio_target)
+        x_start = (w - new_w) // 2
+        image = image.crop((x_start, 0, x_start + new_w, h))
+    else:
+        new_h = int(w / aspect_ratio_target)
+        y_start = (h - new_h) // 2
+        image = image.crop((0, y_start, w, y_start + new_h))
+    return np.array(image.resize((target_w, target_h), PIL.Image.Resampling.BILINEAR))
 
 
 # =========================
-# 🔧 第 5 步：4bit 量化（bitsandbytes NF4）
-#    注意：URSA 的 logits 是用 nn.functional.linear(x, self.lm_head.weight) 算的，
-#    直接吃 weight 张量，所以 lm_head 绝对不能换成 Linear4bit，否则会炸。
+# ✅ 第 6 步：模型加载放在【模块顶层】（ZeroGPU 硬性要求）
 # =========================
-def _linear_targets(root, skip_last_names=("lm_head",)):
-    """收集需要量化的 nn.Linear（按名字最后一节跳过 lm_head 等）。"""
-    out = []
-    for name, mod in root.named_modules():
-        if isinstance(mod, torch.nn.Linear) and name.split(".")[-1] not in skip_last_names:
-            out.append((name, mod))
-    return out
-
-
-def quantize_4bit(root, compute_dtype, quant_type="nf4", blocksize=64, skip_last_names=("lm_head",)):
-    """把 root 下的 nn.Linear 就地换成 bnb Linear4bit，返回 (层数, 原始字节, 量化后字节)。"""
-    import bitsandbytes as bnb
-    from bitsandbytes.functional import quantize_4bit as _q4
-    from bitsandbytes.nn import Linear4bit, Params4bit
-
-    targets = _linear_targets(root, skip_last_names)
-    n, raw_bytes, q_bytes = 0, 0, 0
-    for name, mod in targets:
-        if isinstance(mod, Linear4bit):
-            continue
-        parent = root.get_submodule(name.rsplit(".", 1)[0]) if "." in name else root
-        leaf = name.rsplit(".", 1)[-1]
-
-        w = mod.weight.detach()
-        if w.dtype not in (torch.float32, torch.bfloat16, torch.float16):
-            continue
-        raw_bytes += w.numel() * 4
-        qw, qs = _q4(
-            w.to(torch.float32).contiguous(),
-            blocksize=blocksize,
-            compress_statistics=True,
-            quant_type=quant_type,
-        )
-        new = Linear4bit(
-            mod.in_features,
-            mod.out_features,
-            bias=mod.bias is not None,
-            compute_dtype=compute_dtype,
-            compress_statistics=True,
-            quant_type=quant_type,
-            device=qw.device,
-        )
-        new.weight = Params4bit(
-            qw,
-            requires_grad=False,
-            quant_state=qs,
-            blocksize=blocksize,
-            compress_statistics=True,
-            quant_type=quant_type,
-            quant_storage=qw.dtype,
-            module=new,
-            bnb_quantized=True,
-        )
-        if mod.bias is not None:
-            new.bias.data.copy_(mod.bias.detach())
-        setattr(parent, leaf, new)
-        q_bytes += qw.numel() * qw.element_size() + (qs.absmax.numel() * qs.absmax.element_size())
-        n += 1
-    del targets
-    return n, raw_bytes, q_bytes
+args = parse_args()
 
+# ---------------------------------------------------------------------------
+# ⚠️ torchao int4 的 CUDA 内核（aten::_weight_int4pack_mm）硬性要求：
+#      1) 权重与激活必须是 **bfloat16**（Int4TilePackedTo4dTensor 里直接 assert）
+#      2) 计算能力 >= sm_80（Ampere+）。ZeroGPU 的 A100 / H100 / Blackwell 都满足
+#    所以只要开量化，就强制把精度抬到 bfloat16，避免 fp16 下 dispatch 失败。
+# ---------------------------------------------------------------------------
+if args.quant != "none" and args.precision.lower() != "bfloat16":
+    print(
+        f"[quant] force precision {args.precision} -> bfloat16 "
+        "(torchao int4 tinygemm requires bf16)",
+        flush=True,
+    )
+    args.precision = "bfloat16"
 
-def quantize_int8_dynamic(root):
-    """bnb 不可用时的退路：PyTorch 原生 int8 动态量化（CPU 上有 fbgemm/onednn kernel）。
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu", args.device)
+model_args = {
+    "torch_dtype": getattr(torch, args.precision.lower()),
+    "trust_remote_code": True,
+}
 
-    注意：动态量化要求 fp32 模型，且 lm_head 会被 nn.functional.linear 直接用，
-    必须先从模块树里真正摘掉再装回。
-    """
-    import torch.ao.quantization as tq
-
-    head = root._modules.pop("lm_head", None)  # 真摘除（赋 None 不会从 _modules 移除）
-    try:
-        tq.quantize_dynamic(root, {torch.nn.Linear}, dtype=torch.qint8)
-    finally:
-        if head is not None:
-            root.lm_head = head
-
-
-def self_test_4bit():
-    """不开模型，先验证 CPU 上的 4bit 通路确实可用。"""
-    import bitsandbytes as bnb
-    from bitsandbytes.nn import Linear4bit
-
-    print(f"[test ] bitsandbytes {bnb.__version__}", flush=True)
-    torch.manual_seed(0)
-    m, n, k = 512, 512, 64
-    ref = torch.nn.Linear(k, n, bias=False).eval()
-    q = Linear4bit(k, n, bias=False, compute_dtype=torch.float32, quant_type="nf4", device="cpu")
-    from bitsandbytes.functional import quantize_4bit as _q4
-
-    qw, qs = _q4(ref.weight.detach().contiguous(), compress_statistics=True, quant_type="nf4")
-    from bitsandbytes.nn import Params4bit
-
-    q.weight = Params4bit(qw, requires_grad=False, quant_state=qs, quant_type="nf4",
-                          quant_storage=qw.dtype, module=q, bnb_quantized=True)
-    x = torch.randn(m, k)
-    with torch.no_grad():
-        t0, out = time.time(), q(x)
-        t_q = time.time() - t0
-        t0 = time.time()
-        ref_out = ref(x)
-        t_fp = time.time() - t0
-    err = (out - ref_out).abs().mean().item()
-    print(f"[test ] nf4 err(mean|Δ|)={err:.4f}  4bit {t_q * 1e3:.1f}ms vs fp32 {t_fp * 1e3:.1f}ms "
-          f"({t_fp / max(t_q, 1e-9):.2f}x)", flush=True)
-    return err
+pipe = URSAPipeline.from_pretrained(args.model, **model_args).to(device)
 
 
 # =========================
-# ✅ 第 6 步：参数与模型加载
+# ✅ 第 6.5 步：torchao 权重量化
+#
+#    ⚠️ 为什么仍然不能挪到顶层 CPU 上做？
+#    int4 的 tile_packed_to_4d 布局是 **device-specific** 的：
+#       - 量化时会调 torch.cuda.get_device_capability() 做 assert；
+#       - 打包进 tinygemm 专用 4D 布局（[N/8][K/128][32][4]）；
+#       - transformers 文档也明确 "int4 model can only be loaded on the same
+#         device it was quantized on"。
+#    ZeroGPU 在模块加载阶段 cuda 不可见，所以量化必须【惰性】到 @spaces.GPU 内执行。
+#
+#    ⚠️ packing format 选型（踩坑点）：
+#       - tile_packed_to_4d + hqq  → aten 原生 _weight_int4pack_mm，无额外依赖 ✅ 默认
+#       - plain (Int4Tensor)       → 走 torch.ops.mslk.*，需要 MSLK 可选依赖，
+#                                    没装 mslk 就会在 forward 时炸 ❌
+#       - plain_int32              → XPU / NPU 专用，CUDA 直接不支持 ❌
+#
+#    ⚠️ 量化范围：
+#    - 只替换 nn.Linear：主干 attn(q/k/v/o) + mlp(gate/up/down) 被压缩；
+#    - nn.Embedding（vocab 215669 的 token 表）不动 → 它是显存大头且量化后
+#      forward 要整表反量化，只省存储不省峰值显存；
+#    - lm_head（hidden -> 64000 codebook logits）默认不动，它是采样 token 的最后
+#      一道精度关口，量化后最容易掉画质；
+#    - VAE 默认不动（视频解码直接决定观感）；
+#    - 首尾若干层默认保持 bf16，避免误差在网络两端被放大。
 # =========================
-def parse_args():
-    p = argparse.ArgumentParser(description="URSA TI2V on CPU with 4-bit weights")
-    p.add_argument("--model", default="BAAI/URSA-0.6B-FSQ320")
-    p.add_argument("--quant", default="nf4", choices=["nf4", "fp4", "int8", "none"])
-    p.add_argument("--load-dtype", default="auto", choices=["auto", "bfloat16", "float32"])
-    p.add_argument("--compute-dtype", default="auto", choices=["auto", "bfloat16", "float32"])
-    p.add_argument("--quant-vae", action="store_true", help="连 VAE 里的 Linear 一起量化（画质会掉）")
-    p.add_argument("--threads", type=int, default=0)
-    p.add_argument("--self-test", action="store_true", help="只验证 CPU 4bit 通路，不加载模型")
-    # 生成参数（CLI 模式 / UI 默认值）
-    p.add_argument("--width", type=int, default=256)
-    p.add_argument("--height", type=int, default=160)
-    p.add_argument("--num-frames", type=int, default=9)
-    p.add_argument("--steps", type=int, default=10)
-    p.add_argument("--guidance", type=float, default=7.0)
-    p.add_argument("--motion", type=float, default=9.0)
-    p.add_argument("--seed", type=int, default=0)
-    p.add_argument("--no-randomize", action="store_true")
-    p.add_argument("--prompt", default="")
-    p.add_argument("--negative-prompt", default="worst quality, low quality, inconsistent motion, blurry, jittery")
-    p.add_argument("--image", default="", help="可选：图生视频的首帧图片路径")
-    p.add_argument("--out", default="ursa_out.mp4")
-    p.add_argument("--cli", action="store_true", help="不启动 Gradio，直接生成一次并保存")
-    p.add_argument("--share", action="store_true")
-    p.add_argument("--server-port", type=int, default=7860)
-    return p.parse_args()
+QUANT_MODE = args.quant
+COMPUTE_DTYPE = getattr(torch, args.precision.lower())
+_QUANT_STATE = {"log": "pending"}
 
 
-args = parse_args()
-if args.threads > 0:
-    torch.set_num_threads(args.threads)
-
-if args.self_test:
-    self_test_4bit()
-    raise SystemExit(0)
-
-if args.load_dtype == "auto":
-    args.load_dtype = "bfloat16" if cpu_supports_bf16() else "float32"
-if args.compute_dtype == "auto":
-    args.compute_dtype = args.load_dtype
-load_dtype = getattr(torch, args.load_dtype)
-compute_dtype = getattr(torch, args.compute_dtype)
-
-DEVICE = torch.device("cpu")
-print(f"[load ] model={args.model} quant={args.quant} load_dtype={args.load_dtype} "
-      f"compute_dtype={args.compute_dtype}", flush=True)
-
-pipe = URSAPipeline.from_pretrained(args.model, torch_dtype=load_dtype, trust_remote_code=True)
-pipe.to(DEVICE)
-# 推理路径不会设置 offsets，这里显式清一次，确保走 SDPA 而不是 flex_attention
-try:
-    pipe.transformer.model.flex_attn.offsets = None
-except Exception:
-    pass
-
-# 量化
-if args.quant in ("nf4", "fp4"):
-    try:
-        n, raw_b, q_b = quantize_4bit(pipe.transformer, compute_dtype, quant_type=args.quant)
-        print(f"[quant] transformer: {n} Linear -> {args.quant.upper()} "
-              f"({raw_b / 2**20:.0f}MiB -> {q_b / 2**20:.0f}MiB)", flush=True)
-        if args.quant_vae and hasattr(pipe, "vae"):
-            n2, _, _ = quantize_4bit(pipe.vae, compute_dtype, quant_type=args.quant)
-            print(f"[quant] vae: {n2} Linear -> {args.quant.upper()}", flush=True)
-    except Exception as e:  # bnb 不可用就退回 int8 / 不量化
-        print(f"[quant] bitsandbytes 4bit 失败（{e!r}），回退 int8 动态量化", flush=True)
-        args.quant = "int8"
-if args.quant == "int8":
+def probe_torchao():
+    """探测 torchao 是否可用。构建期无 GPU 也能 import，这里只做可读诊断。"""
+    info = {"available": False, "version": None, "error": None}
     try:
-        if args.load_dtype != "float32":  # 动态量化只吃 fp32
-            print("[quant] int8 动态量化需要 fp32，先把 transformer 转成 float32", flush=True)
-            pipe.transformer.float()
-        quantize_int8_dynamic(pipe.transformer)
-        print("[quant] transformer: int8 dynamic", flush=True)
+        import torchao
+
+        info["available"] = True
+        info["version"] = getattr(torchao, "__version__", "?")
     except Exception as e:
-        print(f"[quant] int8 也失败（{e!r}），以 {args.load_dtype} 直接跑", flush=True)
+        info["error"] = repr(e)
+    print(
+        "[torchao] available=%s version=%s torch=%s cuda=%s quant=%s%s"
+        % (
+            info["available"],
+            info["version"],
+            torch.__version__,
+            torch.version.cuda,
+            QUANT_MODE,
+            " error=%s" % info["error"] if info["error"] else "",
+        ),
+        flush=True,
+    )
+    return info
 
-n_param = sum(p.numel() for p in pipe.transformer.parameters())
-print(f"[load ] transformer params={n_param / 1e6:.1f}M  temporal_stride={pipe.vae_temporal_stride} "
-      f"spatial_stride={pipe.vae_spatial_stride}", flush=True)
 
+_TORCHAO_PROBE = probe_torchao() if QUANT_MODE != "none" else {"available": False}
 
-def latent_tokens(h, w, num_frames):
-    return ((num_frames - 1) // pipe.vae_temporal_stride + 1) * (h // pipe.vae_spatial_stride) * (w // pipe.vae_spatial_stride)
 
+def _is_torchao_weight(w) -> bool:
+    """量化后权重是 torchao 的 Tensor subclass（如 Int4TilePackedTo4dTensor）。"""
+    return isinstance(w, torch.Tensor) and type(w).__module__.startswith("torchao")
 
-# =========================
-# ✅ 第 7 步：推理（无 @spaces.GPU，CPU 上就是普通函数）
-# =========================
-def crop_image(image, target_h, target_w):
-    h, w = image.height, image.width
-    if w / h > target_w / target_h:
-        new_w = int(h * target_w / target_h)
-        image = image.crop(((w - new_w) // 2, 0, (w - new_w) // 2 + new_w, h))
-    else:
-        new_h = int(w * target_h / target_w)
-        image = image.crop((0, (h - new_h) // 2, w, (h - new_h) // 2 + new_h))
-    return np.array(image.resize((target_w, target_h), PIL.Image.Resampling.BILINEAR))
 
+def is_quantized() -> bool:
+    """判断权重是否已量化（fork worker 里也能正确识别）。"""
+    if QUANT_MODE == "none":
+        return True
+    for m in pipe.transformer.modules():
+        if isinstance(m, torch.nn.Linear) and _is_torchao_weight(m.weight):
+            return True
+    if args.quant_vae:
+        for m in pipe.vae.modules():
+            if isinstance(m, torch.nn.Linear) and _is_torchao_weight(m.weight):
+                return True
+    return False
 
-def generate_image(prompt, negative_prompt, seed, randomize_seed, guidance_scale,
-                   size="256x160", num_inference_steps=12):
-    t0 = time.time()
-    w, h = (int(x) for x in size.split("x"))
-    seed = int(np.random.randint(2147483647)) if randomize_seed else int(seed)
-    generator = torch.Generator(device=DEVICE).manual_seed(seed)
-    with torch.no_grad():
-        images = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=h,
-            width=w,
-            num_frames=1,
-            num_inference_steps=int(num_inference_steps),
-            guidance_scale=guidance_scale,
-            generator=generator,
-            output_type="np",
-        ).frames
-    print(f"[t2i  ] {time.time() - t0:.1f}s  tokens={latent_tokens(h, w, 1)}", flush=True)
-    return [export_to_image(image, quality=95) for image in images] + [seed]
 
+def _layer_index(path: str):
+    m = re.search(r"model\.layers\.(\d+)\.", path)
+    return int(m.group(1)) if m else None
 
-def generate_video(prompt, negative_prompt, image, motion_score, seed, randomize_seed,
-                   guidance_scale, num_inference_steps, size="256x160", num_frames=9,
-                   output_type="np"):
-    t0 = time.time()
-    w, h = (int(x) for x in size.split("x"))
-    num_frames, num_inference_steps = int(num_frames), int(num_inference_steps)
-    seed = int(np.random.randint(2147483647)) if randomize_seed else int(seed)
-    generator = torch.Generator(device=DEVICE).manual_seed(seed)
-    with torch.no_grad():
-        frames = pipe(
-            prompt=f"motion={float(motion_score):.1f}, {prompt}",
-            negative_prompt=negative_prompt,
-            image=crop_image(image, h, w) if image is not None else None,
-            height=h,
-            width=w,
-            num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-            output_type=output_type,
-        ).frames[0]
-    dt = time.time() - t0
-    print(f"[ti2v ] {dt:.1f}s  {w}x{h}x{num_frames}  steps={num_inference_steps}  "
-          f"tokens/步≈{latent_tokens(h, w, num_frames) * (2 if guidance_scale > 1 else 1)}", flush=True)
-    return export_to_video(frames, fps=12), seed
 
+def _should_quantize(path: str, n_layers: int) -> bool:
+    if "lm_head" in path and not args.quant_lm_head:
+        return False
+    idx = _layer_index(path)
+    if idx is not None and args.quant_keep_layers > 0 and n_layers > 2 * args.quant_keep_layers:
+        if idx < args.quant_keep_layers or idx >= n_layers - args.quant_keep_layers:
+            return False
+    return True
 
-# =========================
-# ✅ 第 8 步：CLI 模式
-# =========================
-if args.cli or args.prompt:
-    img = PIL.Image.open(args.image) if args.image else None
-    path, used_seed = generate_video(
-        args.prompt or "a lone grizzly bear walks through a misty forest at dawn.",
-        args.negative_prompt,
-        img,
-        args.motion,
-        args.seed,
-        not args.no_randomize,
-        args.guidance,
-        args.steps,
-        size=f"{args.width}x{args.height}",
-        num_frames=args.num_frames,
+
+def build_quant_config():
+    """构造 torchao 的量化 config。"""
+    from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig
+
+    if QUANT_MODE == "int8":
+        return Int8WeightOnlyConfig(group_size=args.quant_group_size, set_inductor_config=False)
+
+    # int4：CUDA 上唯一走 aten 原生 _weight_int4pack_mm 的组合
+    if args.quant_algo == "hqq":
+        return Int4WeightOnlyConfig(
+            group_size=args.quant_group_size,
+            int4_packing_format="tile_packed_to_4d",
+            int4_choose_qparams_algorithm="hqq",
+            set_inductor_config=False,
+        )
+    # tinygemm：量化快（非迭代），但 PLAIN 布局需要 mslk 可选依赖
+    return Int4WeightOnlyConfig(
+        group_size=args.quant_group_size,
+        int4_packing_format="plain",
+        int4_choose_qparams_algorithm="tinygemm",
+        set_inductor_config=False,
+    )
+
+
+@torch.no_grad()
+def apply_torchao_quant_() -> tuple:
+    """用 torchao 原地量化 transformer（可选 VAE）。必须在 CUDA 上执行。"""
+    from torchao.quantization import quantize_
+
+    config = build_quant_config()
+    dev = torch.device("cuda", args.device)
+    n_layers = len(pipe.transformer.model.layers)
+    seen = []
+
+    def filter_fn(module: torch.nn.Module, fqn: str) -> bool:
+        if not isinstance(module, torch.nn.Linear):
+            return False
+        ok = _should_quantize(fqn, n_layers)
+        seen.append((fqn, ok))
+        return ok
+
+    quantize_(pipe.transformer, config, filter_fn=filter_fn, device=dev)
+    quantize_(pipe.vae, config, filter_fn=lambda m, fqn: isinstance(m, torch.nn.Linear)) if (
+        args.quant_vae
+    ) else None
+
+    converted = sum(1 for _, ok in seen if ok)
+    skipped = [fqn for fqn, ok in seen if not ok]
+    return converted, skipped
+
+
+def ensure_model_ready():
+    """把 pipeline 搬到真卡并完成一次性量化。必须在 @spaces.GPU 内部调用。"""
+    global QUANT_MODE
+    dev = torch.device("cuda", args.device)
+    if str(pipe.device) != str(dev):
+        pipe.to(dev)
+    if QUANT_MODE == "none" or is_quantized():
+        return
+    t0 = time.time()
+    try:
+        converted, skipped = apply_torchao_quant_()
+    except Exception as e:
+        # torchao 缺失 / 内核不匹配 / OOM → 优雅降级到原精度，不能让整个生成失败
+        QUANT_MODE = "none"
+        _QUANT_STATE["log"] = f"FALLBACK to {args.precision}: {e!r}"
+        print("[quant]", _QUANT_STATE["log"], flush=True)
+        return
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        mem = torch.cuda.memory_allocated() / 2**30
+    else:  # pragma: no cover
+        mem = -1
+    _QUANT_STATE["log"] = (
+        f"{args.quant}/{args.quant_algo}/g{args.quant_group_size}: {converted} linear converted, "
+        f"{len(skipped)} skipped ({', '.join(skipped[:6])}{' ...' if len(skipped) > 6 else ''}), "
+        f"{time.time() - t0:.1f}s, mem={mem:.2f}GiB"
     )
-    shutil.copy(path, args.out)
-    print(f"[done ] seed={used_seed} -> {os.path.abspath(args.out)}", flush=True)
-    raise SystemExit(0)
+    print("[quant]", _QUANT_STATE["log"], flush=True)
 
 
 # =========================
-# ✅ 第 9 步：Gradio UI
+# ✅ 第 7 步：真正跑推理的函数加 @spaces.GPU
+#    ZeroGPU 对单次 duration 有两道独立检查：
+#      1) 档位上限：超过 → 「illegal duration」（与剩余配额无关，直接拒绝）
+#      2) 配额预扣：声明值 > 访客剩余配额 → 「quota exceeded」
+#    两道都要求声明值越小越好，所以用【动态 duration】按实际步数估算。
 # =========================
-import gradio as gr  # noqa: E402
+
+# 非 PRO 档位的单次上限。若仍报 illegal duration，逐级下调到 60 / 30 再试。
+MAX_DURATION = 120
+
+# 实测基准（half Blackwell, 512x320）：T2I 25 步约 4s，其中首步含预热约 1.3s
+IMAGE_OVERHEAD, IMAGE_PER_STEP = 15, 0.5
+VIDEO_OVERHEAD, VIDEO_PER_STEP = 20, 2.0
+
+# 首次调用额外开销：CPU→GPU 搬运 + int4 打包（HQQ 是迭代算法，group_size 越小越慢）。
+# ZeroGPU 每任务新起进程，量化状态可能不复用，所以按「是否已量化」动态加。
+QUANT_EXTRA_IMAGE, QUANT_EXTRA_VIDEO = 25, 40
+
+
+def _clamped(value):
+    return int(max(20, min(value, MAX_DURATION)))
+
+
+def _quant_extra(kind: str) -> float:
+    if QUANT_MODE == "none" or is_quantized():
+        return 0.0
+    return QUANT_EXTRA_IMAGE if kind == "image" else QUANT_EXTRA_VIDEO
+
+
+def estimate_image_duration(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    num_inference_steps=25,
+):
+    return _clamped(IMAGE_OVERHEAD + num_inference_steps * IMAGE_PER_STEP + _quant_extra("image"))
+
+
+def estimate_video_duration(
+    prompt,
+    negative_prompt,
+    image,
+    motion_score,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    num_inference_steps,
+    output_type="np",
+):
+    return _clamped(VIDEO_OVERHEAD + num_inference_steps * VIDEO_PER_STEP + _quant_extra("video"))
+
+
+@spaces.GPU(duration=estimate_image_duration)
+def generate_image(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    num_inference_steps=25,
+):
+    ensure_model_ready()
+    args = {**locals(), **video_presets["t2i"]}
+    seed = np.random.randint(2147483647) if randomize_seed else seed
+    device = getattr(pipe, "_offload_device", pipe.device)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    images = pipe(generator=generator, **args).frames
+    return [export_to_image(image, quality=95) for image in images] + [seed]
+
+
+@spaces.GPU(duration=estimate_video_duration)
+def generate_video(
+    prompt,
+    negative_prompt,
+    image,
+    motion_score,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    num_inference_steps,
+    output_type="np",
+):
+    ensure_model_ready()
+    args = {**locals(), **video_presets["ti2v"]}
+    args["prompt"] = f"motion={motion_score:.1f}, {prompt}"
+    args["image"] = crop_image(image, args["height"], args["width"]) if image else None
+    seed = np.random.randint(2147483647) if randomize_seed else seed
+    device = getattr(pipe, "_offload_device", pipe.device)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    frames = pipe(generator=generator, **args).frames[0]
+    return export_to_video(frames, fps=12), seed
+
 
 css = """#col-container {margin: 0 auto; max-width: 1366px}"""
 header = (
@@ -503,51 +497,67 @@ header = (
     "<h2>Uniform Discrete Diffusion with Metric Path for Video Generation</h2>"
     "<h3><a href='https://arxiv.org/abs/2510.24717' target='_blank' rel='noopener'>[paper]</a>"
     "<a href='https://github.com/baaivision/URSA' target='_blank' rel='noopener'>[code]</a></h3>"
-    "<p>CPU + 4bit (NF4) 推理：建议 256x160 / 9 帧 / 10 步起步，出片约数分钟。</p>"
+    f"<p style='opacity:.7'>torchao weight-only quantization: <b>{QUANT_MODE}</b>"
+    f" ({args.quant_algo}, group_size={args.quant_group_size}, {args.precision}, "
+    f"keep {args.quant_keep_layers} head/tail layers)</p>"
     "</div>"
 )
 
-SIZE_CHOICES = ["256x160", "320x192", "384x256", "512x320"]
+video_presets = {
+    "t2i": {"width": 512, "height": 320, "num_frames": 1},
+    "ti2v": {"width": 512, "height": 320, "num_frames": 49},
+}
+
 prompts = [
     "a lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur.",
     "Many spotted jellyfish pulsating under water. Their bodies are transparent and glowing in deep ocean.",
-    "An intense close-up of a soldier's face, covered in dirt and sweat, his eyes filled with determination.",
-    "a close-up shot of a woman standing in a dimly lit room, wearing a traditional red and gold chinese outfit.",
-    "The camera slowly rotates around a massive stack of vintage televisions in a large museum gallery.",
+    "An intense close-up of a soldier’s face, covered in dirt and sweat, his eyes filled with determination as he surveys the battlefield.",
+    "a close-up shot of a woman standing in a dimly lit room. she is wearing a traditional chinese outfit, which includes a red and gold dress with intricate designs and a matching headpiece. the woman has her hair styled in an updo, adorned with a gold accessory. her makeup is done in a way that accentuates her features, with red lipstick and dark eyeshadow. she is looking directly at the camera with a neutral expression. the room has a rustic feel, with wooden beams and a stone wall visible in the background. the lighting in the room is soft and warm, creating a contrast with the woman's vibrant attire. there are no texts or other objects in the video. the style of the video is a portrait, focusing on the woman and her attire.",
+    "The camera slowly rotates around a massive stack of vintage televisions that are placed within a large New York museum gallery. Each of the televisions is showing a different program. There are 1950s sci-fi movies with their distinctive visuals, horror movies with their creepy scenes, news broadcasts with moving images and words, static on some screens, and a 1970s sitcom with its characteristic look. The televisions are of various sizes and designs, some with rounded edges and others with more angular shapes. The gallery is well-lit, with light falling on the stack of televisions and highlighting the different programs being shown. There are no people visible in the immediate vicinity, only the stack of televisions and the surrounding gallery space.",
 ]
-examples = [[p, 9] for p in prompts]
+motion_scores = [9, 9, 9, 9, 9]
+examples = [list(x) for x in zip(prompts, motion_scores)]
+
 
 if __name__ == "__main__":
+    # Application
     app = gr.Blocks(css=css, theme="origin").__enter__()
     container = gr.Column(elem_id="col-container").__enter__()
     _, main_row = gr.Markdown(header), gr.Row().__enter__()
 
+    # Input
     input_col = gr.Column().__enter__()
-    prompt = gr.Text(label="Prompt", lines=5, value=prompts[0])
+    prompt = gr.Text(
+        label="Prompt",
+        placeholder="Describe the video you want to generate",
+        value="A lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur.",
+        lines=5,
+    )
     negative_prompt = gr.Text(
-        label="Negative Prompt", lines=1,
-        value="worst quality, low quality, inconsistent motion, static, still, blurry, jittery, ugly",
+        label="Negative Prompt",
+        placeholder="Describe what you don't want in the video",
+        value="worst quality, low quality, inconsistent motion, static, still, blurry, jittery, distorted, ugly",
+        lines=1,
     )
     with gr.Row():
         generate_image_btn = gr.Button("Generate Image Prompt", variant="primary", size="lg")
         generate_video_btn = gr.Button("Generate Video", variant="primary", size="lg")
-    image_prompt = gr.Image(label="Image Prompt", height=320, type="pil")
+    image_prompt = gr.Image(label="Image Prompt", height=480, type="pil")
 
-    options = gr.Accordion("Options", open=True).__enter__()
-    with gr.Row():
-        size = gr.Dropdown(label="Resolution", choices=SIZE_CHOICES, value="256x160")
-        num_frames = gr.Dropdown(label="Frames", choices=[1, 5, 9, 17, 25, 49], value=9)
-    with gr.Row():
-        seed = gr.Slider(label="Seed", maximum=2147483647, step=1, value=0)
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+    options = gr.Accordion("Options", open=False).__enter__()
+    seed = gr.Slider(label="Seed", maximum=2147483647, step=1, value=0)
+    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+    guidance_scale = gr.Slider(label="Guidance scale", minimum=1, maximum=10.0, step=0.1, value=7.0)
     with gr.Row():
-        guidance_scale = gr.Slider(label="Guidance scale", minimum=1, maximum=10.0, step=0.1, value=7.0)
-        num_inference_steps = gr.Slider(label="Inference steps", minimum=2, maximum=30, step=1, value=10)
+        # 上限从 100 收到 50：与 estimate_video_duration 的估算保持自洽，
+        # 避免访客拉满步数后必然撞上档位上限而被直接拒绝
+        num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=30)
     options.__exit__(), input_col.__exit__()
 
+    # Results
     result_col = gr.Column().__enter__()
     motion = gr.Slider(label="Motion Score", minimum=1, maximum=10, step=1, value=9)
-    result = gr.Video(label="Result", height=320, show_label=False, autoplay=True)
+    result = gr.Video(label="Result", height=480, show_label=False, autoplay=True)
     result_col.__exit__(), main_row.__exit__()
 
     with gr.Row():
@@ -555,18 +565,33 @@ if __name__ == "__main__":
 
     container.__exit__()
     gr.on(
-        triggers=[generate_image_btn.click],
+        triggers=[generate_image_btn.click, prompt.submit, negative_prompt.submit],
         fn=generate_image,
-        inputs=[prompt, negative_prompt, seed, randomize_seed, guidance_scale, size, num_inference_steps],
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            guidance_scale,
+        ],
         outputs=[image_prompt, seed],
     )
     gr.on(
-        triggers=[generate_video_btn.click],
+        triggers=[generate_video_btn.click, prompt.submit, negative_prompt.submit],
         fn=generate_video,
-        inputs=[prompt, negative_prompt, image_prompt, motion, seed, randomize_seed,
-                guidance_scale, num_inference_steps, size, num_frames],
+        inputs=[
+            prompt,
+            negative_prompt,
+            image_prompt,
+            motion,
+            seed,
+            randomize_seed,
+            guidance_scale,
+            num_inference_steps,
+        ],
         outputs=[result, seed],
     )
+    # 限制并发，避免多个请求抢同一张卡
     app.__exit__(), app.queue(default_concurrency_limit=1).launch(
-        share=args.share, ssr_mode=False, max_threads=1, server_port=args.server_port
+        share=False, ssr_mode=False, max_threads=1
     )