Load model directly

app.py

@@ -34,26 +34,28 @@ def load_model():
         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]
+    # # 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]
+    # 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}")
+    # 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)
+    # 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, weights_only=True)
+    model = torch.load(checkpoint_path, map_location="cuda")
+
 
     # 4. Final setup
     model = disable_dropout(model)