Load model differently

app.py

@@ -4,8 +4,10 @@ import numpy as np
 import json
 import time
 from transformers import AutoTokenizer
-from llama_diffusion_model import CustomTransformerModel, CustomTransformerConfig, disable_dropout
+from llama_diffusion_model import disable_dropout
 import os
+import importlib
+from huggingface_hub import hf_hub_download
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -24,18 +26,43 @@ token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(toke
 
 
 def load_model():
-    config = CustomTransformerConfig(vocab_size=vocab_size)
-    model = CustomTransformerModel(config)
-    model.load_state_dict(torch.hub.load_state_dict_from_url(
-        "https://huggingface.co/Ruurd/tini_model/resolve/main/diffusion-model.pth",
-        map_location="cuda",
-        headers={"Authorization": f"Bearer {hf_token}"}
-    ))
+
+    # 1. Download the checkpoint
+    checkpoint_path = hf_hub_download(
+        repo_id="ruurd/diffusion-llama",
+        filename="diffusion-model.pth",
+        token=os.getenv("HF_TOKEN")
+    )
+
+    # 2. Prepare dynamic class loading like you did before
+    torch.serialization.clear_safe_globals()
+    unsafe_globals = torch.serialization.get_unsafe_globals_in_checkpoint(checkpoint_path)
+    missing_class_names = [name.split(".")[-1] for name in unsafe_globals]
+
+    safe_classes = [cls for name, cls in globals().items() if name in missing_class_names]
+
+    for class_path in unsafe_globals:
+        try:
+            module_name, class_name = class_path.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            cls = getattr(module, class_name)
+            safe_classes.append(cls)
+        except (ImportError, AttributeError) as e:
+            print(f"⚠️ Warning: Could not import {class_path} - {e}")
+
+    torch.serialization.add_safe_globals(safe_classes)
+
+    # 3. Actually load the full model
+    model = torch.load(checkpoint_path, weights_only=True)
+
+    # 4. Final setup
     model = disable_dropout(model)
     model.to("cuda")
     model.eval()
+
     return model
 
+
 rng = np.random.default_rng()
 
 # --- Utility Functions ---