Update gen test models

Author: qubvel

misc/generate_test_models.py

@@ -9,33 +9,50 @@
 
 api = huggingface_hub.HfApi(token=os.getenv("HF_TOKEN"))
 
-for model_name, model_class in smp.MODEL_ARCHITECTURES_MAPPING.items():
-    model = model_class(encoder_name=ENCODER_NAME)
-    model = model.eval()
-
-    # generate test sample
-    torch.manual_seed(423553)
-    sample = torch.rand(1, 3, 256, 256)
-
-    with torch.no_grad():
-        output = model(sample)
 
+def save_and_push(model, inputs, outputs, model_name, encoder_name):
     with tempfile.TemporaryDirectory() as tmpdir:
         # save model
         model.save_pretrained(f"{tmpdir}")
 
         # save input and output
-        torch.save(sample, f"{tmpdir}/input-tensor.pth")
-        torch.save(output, f"{tmpdir}/output-tensor.pth")
+        torch.save(inputs, f"{tmpdir}/input-tensor.pth")
+        torch.save(outputs, f"{tmpdir}/output-tensor.pth")
 
         # create repo
-        repo_id = f"{HUB_REPO}/{model_name}-{ENCODER_NAME}"
+        repo_id = f"{HUB_REPO}/{model_name}-{encoder_name}"
         if not api.repo_exists(repo_id=repo_id):
             api.create_repo(repo_id=repo_id, repo_type="model")
 
         # upload to hub
         api.upload_folder(
             folder_path=tmpdir,
-            repo_id=f"{HUB_REPO}/{model_name}-{ENCODER_NAME}",
+            repo_id=f"{HUB_REPO}/{model_name}-{encoder_name}",
             repo_type="model",
         )
+
+
+for model_name, model_class in smp.MODEL_ARCHITECTURES_MAPPING.items():
+    if model_name == "dpt":
+        encoder_name = "tu-test_vit"
+        model = smp.DPT(
+            encoder_name=encoder_name,
+            decoder_readout="cat",
+            decoder_intermediate_channels=(16, 32, 64, 64),
+            decoder_fusion_channels=16,
+            dynamic_img_size=True,
+        )
+    else:
+        encoder_name = ENCODER_NAME
+        model = model_class(encoder_name=encoder_name)
+
+    model = model.eval()
+
+    # generate test sample
+    torch.manual_seed(423553)
+    sample = torch.rand(1, 3, 256, 256)
+
+    with torch.no_grad():
+        output = model(sample)
+
+    save_and_push(model, sample, output, model_name, encoder_name)