Upload backend.py

backend.py

@@ -61,6 +61,57 @@ DEVICE, DTYPE = select_device()
 # YOLO runs outside the @spaces.GPU call, and on ZeroGPU CUDA is only usable inside it
 POSE_DEVICE = "cpu" if os.environ.get("SPACES_ZERO_GPU") else DEVICE.type
 
+
+# ---------- Diagnostics (set SLT_DEBUG=0 to silence) ----------
+DEBUG = os.environ.get("SLT_DEBUG", "1") != "0"
+
+
+def diag(message):
+    if DEBUG:
+        print(f"[diag] {message}", flush=True)
+
+
+def tensor_summary(t):
+    """Shape, dtype, device and value range of a tensor (NaN/Inf flagged)."""
+    if not isinstance(t, torch.Tensor):
+        return repr(t)[:60]
+    f = t.detach().float()
+    finite = f[torch.isfinite(f)]
+    rng = f"min={finite.min().item():.4g} max={finite.max().item():.4g} mean={finite.mean().item():.4g}" if finite.numel() else "no finite values"
+    bad = f" NaN={torch.isnan(f).sum().item()} Inf={torch.isinf(f).sum().item()}" if finite.numel() != f.numel() else ""
+    return f"{tuple(t.shape)} {t.dtype} {t.device} {rng}{bad}"
+
+
+def weights_fingerprint(m):
+    """A few numbers that change if weights are lost, zeroed or garbled when moved between devices."""
+    params = dict(m.named_parameters())
+    names = [next(iter(params)), "mt5_model.shared.weight", "mt5_model.lm_head.weight", list(params)[-1]]
+    return {n.split(".")[-2] + "." + n.split(".")[-1]: round(params[n].detach().float().abs().sum().item(), 3) for n in names if n in params}
+
+
+def weights_location(m):
+    """How many parameters sit on each device/dtype."""
+    counts = {}
+    for p in m.parameters():
+        key = f"{p.device}/{p.dtype}"
+        counts[key] = counts.get(key, 0) + p.numel()
+    return {k: f"{v / 1e6:.1f}M" for k, v in counts.items()}
+
+
+def environment_report():
+    import transformers
+    try:
+        import spaces as _spaces_module  # noqa: F401
+        has_spaces = not isinstance(spaces, type)
+    except ImportError:
+        has_spaces = False
+    cuda = torch.cuda.is_available()
+    diag(f"python {sys.version.split()[0]} | torch {torch.__version__} (cuda build {torch.version.cuda}) | transformers {transformers.__version__}")
+    diag(f"spaces package: {has_spaces} | SPACES_ZERO_GPU={os.environ.get('SPACES_ZERO_GPU')!r} | "
+         f"SLT_DEVICE={os.environ.get('SLT_DEVICE', 'auto')!r} SLT_DTYPE={os.environ.get('SLT_DTYPE', 'auto')!r}")
+    diag(f"torch.cuda.is_available()={cuda}" + (f" | {torch.cuda.device_count()} GPU(s): {torch.cuda.get_device_name(0)} | bf16 supported: {torch.cuda.is_bf16_supported()}" if cuda else ""))
+    diag(f"selected DEVICE={DEVICE} DTYPE={DTYPE} POSE_DEVICE={POSE_DEVICE}")
+
 # Global variables to cache models and config so they are not reloaded on every video translation
 model = None
 pose_models = None
@@ -92,19 +143,28 @@ def initialize_model():
         return
 
     print("Initializing Uni_Sign model...")
+    environment_report()
     args = InferenceConfig()
     model = Uni_Sign(args=args)
 
     # Load model weights
     if args.finetune and os.path.exists(args.finetune):
         print(f"Loading checkpoint weights from: {args.finetune}")
+        diag(f"checkpoint size: {os.path.getsize(args.finetune) / 1e6:.1f} MB")
         state_dict = torch.load(args.finetune, map_location='cpu')['model']
-        model.load_state_dict(state_dict, strict=False)
+        result = model.load_state_dict(state_dict, strict=False)
+        diag(f"checkpoint tensors: {len(state_dict)} | model tensors: {len(model.state_dict())} | "
+             f"missing: {len(result.missing_keys)} {result.missing_keys[:5]} | unexpected: {len(result.unexpected_keys)} {result.unexpected_keys[:5]}")
     else:
         print(f"WARNING: Checkpoint weights not found at '{args.finetune}'!")
 
+    diag(f"weights fingerprint after loading (cpu, float32): {weights_fingerprint(model)}")
     model.to(device=DEVICE, dtype=DTYPE).eval()
     print(f"Uni_Sign model on {DEVICE} ({DTYPE}); pose detection on {POSE_DEVICE}")
+    diag(f"weights location after .to(): {weights_location(model)}")
+    diag(f"weights fingerprint after .to(): {weights_fingerprint(model)}")
+    diag(f"mT5 generation config: decoder_start={model.mt5_model.config.decoder_start_token_id} "
+         f"eos={model.mt5_model.config.eos_token_id} pad={model.mt5_model.config.pad_token_id}")
 
     print("Initializing Mediapipe and YOLO models...")
     pose_checkpoint_folder = 'checkpoints/pose/'
@@ -156,6 +216,10 @@ def process_pose_data_in_memory(pose_results, args):
                  formatted_frame[part] = np.array(kps)[:, :2].tolist()
         pose.append(formatted_frame)
 
+    detected = {part: sum(1 for f in pose if len(f[part]) > 0) for part in YTASL_GROUP_SIZES}
+    diag(f"keypoints: {len(pose)} frames; frames with detections per part: "
+         + ", ".join(f"{part}={n}/{len(pose)}" for part, n in detected.items()))
+
     # Subsample frames to match target length (max_length)
     duration = len(pose)
     tmp = select_frame_indices(duration, args.max_length, phase='test')
@@ -194,21 +258,48 @@ def process_pose_data_in_memory(pose_results, args):
     src_input['src_length_batch'] = torch.LongTensor([seq_len])
     src_input['name_batch'] = ["gradio_video"]
 
+    for key, value in src_input.items():
+        if isinstance(value, torch.Tensor):
+            diag(f"model input {key}: {tensor_summary(value)}")
     return src_input
 
 
 @spaces.GPU(duration=60)  # on ZeroGPU, only this call holds a GPU (time limit in seconds)
 def run_translation(src_input):
     """Runs Uni_Sign on prepared pose tensors and returns the decoded translation."""
+    # Inside the GPU call: is CUDA really there, and did the weights arrive intact?
+    cuda = torch.cuda.is_available()
+    diag(f"in run_translation: torch.cuda.is_available()={cuda}"
+         + (f", current device {torch.cuda.current_device()} ({torch.cuda.get_device_name(0)})" if cuda else ""))
+    diag(f"weights location in run_translation: {weights_location(model)}")
+    diag(f"weights fingerprint in run_translation: {weights_fingerprint(model)}")
+
     for key, value in src_input.items():
         if isinstance(value, torch.Tensor):
             src_input[key] = value.to(device=DEVICE, dtype=DTYPE if value.is_floating_point() else value.dtype)
     tgt_input = {'gt_sentence': [""], 'gt_gloss': [""]}
+    tokenizer = model.mt5_tokenizer
 
     with torch.inference_mode():
         stack_out = model(src_input, tgt_input)
+        diag(f"encoder input embeds: {tensor_summary(stack_out['inputs_embeds'])}")
+        diag(f"encoder attention mask: {tensor_summary(stack_out['attention_mask'])}")
+
+        # What the model wants to say first: top-5 candidates for the first output token
+        start = torch.full((1, 1), model.mt5_model.config.decoder_start_token_id, device=DEVICE, dtype=torch.long)
+        logits = model.mt5_model(inputs_embeds=stack_out['inputs_embeds'], attention_mask=stack_out['attention_mask'],
+                                 decoder_input_ids=start).logits[0, -1].float()
+        diag(f"first-token logits: {tensor_summary(logits)}")
+        probs, ids = logits.softmax(-1).topk(5)
+        diag("first-token top-5: " + ", ".join(f"{tok!r}={p:.3f}" for tok, p in
+                                               zip(tokenizer.convert_ids_to_tokens(ids.tolist()), probs.tolist())))
+
+        greedy = model.generate(stack_out, max_new_tokens=100, num_beams=1)
+        diag(f"greedy: ids {greedy[0].tolist()[:20]} -> {tokenizer.batch_decode(greedy, skip_special_tokens=True)[0]!r}")
         output = model.generate(stack_out, max_new_tokens=100, num_beams=4)
-    return model.mt5_tokenizer.batch_decode(output, skip_special_tokens=True)[0].strip()
+        diag(f"beam search (4): ids {output[0].tolist()[:20]} -> {tokenizer.batch_decode(output, skip_special_tokens=True)[0]!r}")
+
+    return tokenizer.batch_decode(output, skip_special_tokens=True)[0].strip()
 
 
 def process_input(input_video_path, progress=None, log_callback=None):