Update mbart50.py

Physton · Jun 24, 2023 · 6976a4c · 6976a4c
1 parent 1cd5ff3
commit 6976a4c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/scripts/physton_prompt/mbart50.py b/scripts/physton_prompt/mbart50.py
@@ -15,6 +15,12 @@ def initialize(reload=False):
     if not reload and model is not None:
         return
     loading = True
+
+    model_path = os.path.join(cache_dir, "mbart-large-50-many-to-many-mmt")
+    model_file = os.path.join(model_path, "pytorch_model.bin")
+    if os.path.exists(model_path) and os.path.exists(model_file):
+        model_name = model_path
+
     from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
     print(f'[sd-webui-prompt-all-in-one] Loading model {model_name} from {cache_dir}...')
     model = MBartForConditionalGeneration.from_pretrained(model_name, cache_dir=cache_dir)