فهرست منبع

fix jit_export

lyuxiang.lx 4 ماه پیش
والد
کامیت
59cb2bf16c
1فایلهای تغییر یافته به همراه8 افزوده شده و 8 حذف شده
  1. 8 8
      cosyvoice/bin/export_jit.py

+ 8 - 8
cosyvoice/bin/export_jit.py

@@ -62,14 +62,6 @@ def main():
     model = AutoModel(model_dir=args.model_dir)
 
     if get_model_type(model.model) == CosyVoiceModel:
-        # 1. export flow encoder
-        flow_encoder = model.model.flow.encoder
-        script = get_optimized_script(flow_encoder)
-        script.save('{}/flow.encoder.fp32.zip'.format(args.model_dir))
-        script = get_optimized_script(flow_encoder.half())
-        script.save('{}/flow.encoder.fp16.zip'.format(args.model_dir))
-        logging.info('successfully export flow_encoder')
-    elif get_model_type(model.model) == CosyVoice2Model:
         # 1. export llm text_encoder
         llm_text_encoder = model.model.llm.text_encoder
         script = get_optimized_script(llm_text_encoder)
@@ -93,6 +85,14 @@ def main():
         script = get_optimized_script(flow_encoder.half())
         script.save('{}/flow.encoder.fp16.zip'.format(args.model_dir))
         logging.info('successfully export flow_encoder')
+    elif get_model_type(model.model) == CosyVoice2Model:
+        # 1. export flow encoder
+        flow_encoder = model.model.flow.encoder
+        script = get_optimized_script(flow_encoder)
+        script.save('{}/flow.encoder.fp32.zip'.format(args.model_dir))
+        script = get_optimized_script(flow_encoder.half())
+        script.save('{}/flow.encoder.fp16.zip'.format(args.model_dir))
+        logging.info('successfully export flow_encoder')
     else:
         raise ValueError('unsupported model type')