Try again

app.py

@@ -8,6 +8,7 @@ from llama_diffusion_model import disable_dropout
 import os
 import importlib
 from huggingface_hub import hf_hub_download
+from llama_diffusion_model import CustomTransformerModel, CustomTransformerConfig, disable_dropout
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -26,42 +27,16 @@ token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(toke
 
 
 def load_model():
-
-    # 1. Download the checkpoint
-    checkpoint_path = hf_hub_download(
-        repo_id="ruurd/tini_model",
+    ckpt_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)
-    model = torch.load(checkpoint_path, map_location="cuda")
-
-
-    # 4. Final setup
+    model = torch.load(ckpt_path, map_location="cuda")  # no weights_only, no globals hack
     model = disable_dropout(model)
     model.to("cuda")
     model.eval()
-
     return model