interior-ai-designer

Sleeping

Bobby commited on Jul 17, 2024

Commit

8991de8

1 Parent(s): c287f98

scannet

Files changed (1) hide show

app.py CHANGED Viewed

@@ -56,7 +56,7 @@ class Preprocessor:
             return
         elif name == "NormalBae":
             print("Loading NormalBae")
-            model_file = hf_hub_download(self.MODEL_ID, filename="NormalBaeDetector.pth")
             self.model = NormalBaeDetector.from_pretrained(model_file).to("cuda")
             torch.cuda.empty_cache()
             self.name = name

             return
         elif name == "NormalBae":
             print("Loading NormalBae")
+            model_file = hf_hub_download(self.MODEL_ID, filename="scannet.pt")
             self.model = NormalBaeDetector.from_pretrained(model_file).to("cuda")
             torch.cuda.empty_cache()
             self.name = name