Change back to normal 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