Monkey-patch (temporarily) for LoRA layers

app.py

@@ -25,6 +25,20 @@ 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_m",
@@ -34,6 +48,16 @@ def load_model():
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = torch.load(ckpt_path, map_location=device)
+
+    # 🔧 Monkey-patch for missing PEFT attribute
+    def add_cast_input_dtype_enabled(module):
+        for child in module.children():
+            add_cast_input_dtype_enabled(child)
+        if isinstance(module, torch.nn.Linear) and not hasattr(module, "cast_input_dtype_enabled"):
+            module.cast_input_dtype_enabled = False
+
+    add_cast_input_dtype_enabled(model)
+
     model = disable_dropout(model)
     model.to(device)
     model.eval()