feat(ml):debug custom jit

models/diffusion_networks.py

@@ -289,6 +289,16 @@ def define_G(
 
     elif G_netG == "vit_vid":
         variant = getattr(opt, "G_vit_variant", "")
+        if variant and variant not in JiTVid_VARIANT_CONFIGS:
+            if variant.startswith("JiT-"):
+                alias = f"JiTVid-{variant[len('JiT-'):]}"
+                if alias in JiTVid_VARIANT_CONFIGS:
+                    variant = alias
+            if variant not in JiTVid_VARIANT_CONFIGS:
+                raise ValueError(
+                    f"Unknown G_vit_variant '{variant}'. "
+                    f"Valid: {sorted(JiTVid_VARIANT_CONFIGS.keys())}"
+                )
         base = JiTVid_VARIANT_CONFIGS.get(variant, {})
         cfg = {
             "depth": getattr(opt, "G_vit_depth", base.get("depth", 12)),

options/common_options.py

@@ -428,7 +428,7 @@ def initialize(self, parser):
             "--G_vit_variant",
             type=str,
             default="JiT-B/16",
-            help="Selects the ViT backbone when --G_netG vit",
+            help="Selects the ViT backbone when --G_netG vit (use JiT-*) or vit_vid (use JiTVid-*)",
         )
         parser.add_argument(
             "--G_vit_disable_bottleneck",

scripts/gen_single_image_diffusion.py

@@ -81,7 +81,7 @@ def load_model(
     if opt.model_type in ["cm", "cm_gan", "sc", "b2b"]:
         opt.alg_palette_sampling_method = sampling_method
         opt.alg_diffusion_cond_embed_dim = 256
-    model = diffusion_networks.define_G(**vars(opt))
+    model = diffusion_networks.define_G(opt=opt, **vars(opt))
     model.eval()
 
     # handle old models

scripts/gen_vid_autoregressive_diffusion_forward_NoCanny_offline.py

@@ -101,7 +101,7 @@ def load_model(
     if opt.model_type in ["cm", "cm_gan", "b2b"]:
         opt.alg_palette_sampling_method = sampling_method
         opt.alg_diffusion_cond_embed_dim = 256
-    model = diffusion_networks.define_G(**vars(opt))
+    model = diffusion_networks.define_G(opt=opt, **vars(opt))
     model.eval()
 
     # handle old models