Change to mask model

app.py

@@ -25,51 +25,51 @@ with open("token_probabilities.json") as f:
     token_probs_dict = json.load(f)
 token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(token_probs_dict))], dtype=np.float32)
 
-# def load_model():
-#     ckpt_path = hf_hub_download(
-#         repo_id="ruurd/tini_bi_m",
-#         filename="diffusion-model.pth",
-#         token=os.getenv("HF_TOKEN")
-#     )
-
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     model = torch.load(ckpt_path, map_location=device)
-#     model = disable_dropout(model)
-#     model.to(device)
-#     model.eval()
-#     return model
-
 def load_model():
     ckpt_path = hf_hub_download(
-        repo_id="ruurd/tini_bi",
+        repo_id="ruurd/tini_bi_m",
         filename="diffusion-model.pth",
-        token=os.getenv("HF_TOKEN"),
-        revision="5a22a8b6168466dbbf704efd00d8cbf2eee51426",
+        token=os.getenv("HF_TOKEN")
     )
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = torch.load(ckpt_path, map_location=device)
+    model = disable_dropout(model)
+    model.to(device)
+    model.eval()
+    return model
 
-    # Step 1: Create model from scratch
-    model = CustomTransformerModel(CustomTransformerConfig())
+# def load_model():
+#     ckpt_path = hf_hub_download(
+#         repo_id="ruurd/tini_bi",
+#         filename="diffusion-model.pth",
+#         token=os.getenv("HF_TOKEN"),
+#         revision="5a22a8b6168466dbbf704efd00d8cbf2eee51426",
+#     )
 
-    # Step 2: Load state_dict from full checkpoint
-    full_model = torch.load(ckpt_path, map_location=device)
+#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # This handles both full model or just state_dict
-    try:
-        state_dict = full_model.state_dict()
-    except AttributeError:
-        state_dict = full_model  # already a state_dict
+#     # Step 1: Create model from scratch
+#     model = CustomTransformerModel(CustomTransformerConfig())
 
-    # Step 3: Load weights (might print mismatches)
-    missing, unexpected = model.load_state_dict(state_dict, strict=False)
-    print("Missing keys:", missing)
-    print("Unexpected keys:", unexpected)
+#     # Step 2: Load state_dict from full checkpoint
+#     full_model = torch.load(ckpt_path, map_location=device)
 
-    model = disable_dropout(model)
-    model.to(device)
-    model.eval()
-    return model
+#     # This handles both full model or just state_dict
+#     try:
+#         state_dict = full_model.state_dict()
+#     except AttributeError:
+#         state_dict = full_model  # already a state_dict
+
+#     # Step 3: Load weights (might print mismatches)
+#     missing, unexpected = model.load_state_dict(state_dict, strict=False)
+#     print("Missing keys:", missing)
+#     print("Unexpected keys:", unexpected)
+
+#     model = disable_dropout(model)
+#     model.to(device)
+#     model.eval()
+#     return model