Upload infer/common_infer.py with huggingface_hub

infer/common_infer.py

@@ -29,6 +29,14 @@ def resolve_local_config_path(path: str | None) -> str | None:
     return os.path.join(PROJECT_ROOT, path)
 
 
+def load_trusted_checkpoint(path: str, map_location=None):
+    """Load SynLayers-owned checkpoints saved before PyTorch changed weights_only defaults."""
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
 def scale_box_xyxy(box, source_size: int, target_size: int) -> tuple:
     """Scale an xyxy box from source_size to target_size."""
     scale = target_size / source_size
@@ -92,7 +100,9 @@ def initialize_pipeline(config):
         single_layer_decoder=config.get("single_layer_decoder", None),
     )
     transp_vae = CustomVAE(vae_args)
-    transp_vae_weights = torch.load(transp_vae_path, map_location=torch.device("cuda"))
+    transp_vae_weights = load_trusted_checkpoint(
+        transp_vae_path, map_location=torch.device("cuda")
+    )
     missing_keys, unexpected_keys = transp_vae.load_state_dict(
         transp_vae_weights["model"], strict=False
     )
@@ -141,7 +151,7 @@ def initialize_pipeline(config):
     layer_pe_path = os.path.join(layer_ckpt, "layer_pe.pth") if layer_ckpt else ""
     if os.path.exists(layer_pe_path):
         print(f"[INFO] Loading layer_pe from {layer_pe_path}...", flush=True)
-        layer_pe = torch.load(layer_pe_path)
+        layer_pe = load_trusted_checkpoint(layer_pe_path, map_location=torch.device("cuda"))
         transformer.load_state_dict(layer_pe, strict=False)
 
     print("[INFO] Loading MultiLayer-Adapter...", flush=True)
@@ -174,4 +184,4 @@ def initialize_pipeline(config):
         print(f"[INFO] Loading trained LoRA from {lora_ckpt}...", flush=True)
         pipeline.load_lora_weights(lora_ckpt, adapter_name="layer")
 
-    return pipeline, transp_vae
+    return pipeline, transp_vae