fix SD loading when safetensors weights only

huggingface · Jul 27, 2023 · d3ce904 · d3ce904
1 parent ab0df1b
commit d3ce904
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1240,7 +1240,7 @@ def determine_framework(
             framework = "pt"
         elif any(is_tf_weight_file):
             framework = "tf"
-        elif "model_index.json" in all_files and any(file.endswith(Path(WEIGHTS_NAME).suffix) for file in all_files):
+        elif "model_index.json" in all_files and any(file.endswith((pt_weight_extension, safe_weight_extension)) for file in all_files):
             # stable diffusion case
             framework = "pt"
         else: