Skip to content

Commit 06ae2ff

Browse files
committed
fix issues with automatic compress model file in tar format
1 parent 216a63c commit 06ae2ff

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

webui.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -746,15 +746,16 @@ def train():
746746
#automatic tar latest checkpoint and upload to s3 by zheng on 2023.03.22
747747
os.makedirs(os.path.dirname("/opt/ml/model/"), exist_ok=True)
748748
train_steps=int(db_config.revision)
749-
f1=os.path.join(sd_models_path, db_model_name, f'{db_model_name}_{train_steps}.yaml')
749+
model_file_basename = f'{db_model_name}_{train_steps}_lora' if db_config.use_lora else f'{db_model_name}_{train_steps}'
750+
f1=os.path.join(sd_models_path, db_model_name, f'{model_file_basename}.yaml')
750751
if os.path.exists(f1):
751752
shutil.copy(f1,"/opt/ml/model/")
752753
if db_save_safetensors:
753-
f2=os.path.join(sd_models_path, db_model_name, f'{db_model_name}_{train_steps}.safetensors')
754+
f2=os.path.join(sd_models_path, db_model_name, f'{model_file_basename}.safetensors')
754755
if os.path.exists(f2):
755756
shutil.copy(f2,"/opt/ml/model/")
756757
else:
757-
f2=os.path.join(sd_models_path, db_model_name, f'{db_model_name}_{train_steps}.ckpt')
758+
f2=os.path.join(sd_models_path, db_model_name, f'{model_file_basename}.ckpt')
758759
if os.path.exists(f2):
759760
shutil.copy(f2,"/opt/ml/model/")
760761
except Exception as e:

0 commit comments

Comments
 (0)