Upload 18 files

ace_test.py

@@ -0,0 +1,45 @@
+from acestep.pipeline_ace_step import ACEStepPipeline
+from ov_ace_helper import OVACEStepPipeline
+import os
+import requests
+import platform
+from pathlib import Path
+
+inputs = {
+    "prompt": "country rock, folk rock, southern rock, bluegrass, country pop",
+    "lyrics": "[verse]\nWoke up to the sunrise glow\nTook my heart and hit the road[inst]",
+    "audio_duration": 15.0,
+    "infer_step": 25,
+    "use_erg_tag": False,
+    "use_erg_lyric": True,
+    "use_erg_diffusion": True,
+    "save_path": Path("outputs").absolute().as_posix(),
+    "task": "text2music",
+}
+
+if not Path(inputs["save_path"]).exists():
+    os.mkdir(inputs["save_path"])
+
+checkpoint_dir = ""
+pipeline = ACEStepPipeline(checkpoint_dir=checkpoint_dir, dtype="float32", cpu_offload=False)
+pipeline.load_checkpoint(checkpoint_dir)
+
+result = pipeline(**inputs)
+output_path = result[0]
+print(output_path)
+
+import nncf
+from ov_ace_helper import convert_models
+
+ov_converted_model_dir = "ov_models"
+weights_compression_config = {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 128, "ratio": 0.8}
+ov_converted_model_dir += "_int4"
+
+convert_models(pipeline, model_dir=ov_converted_model_dir, orig_checkpoint_path=checkpoint_dir, quantization_config=weights_compression_config)
+
+ov_pipeline = OVACEStepPipeline()
+ov_pipeline.load_models(ov_models_path=ov_converted_model_dir, device='CPU')
+
+ov_result = ov_pipeline(**inputs)
+
+ov_out_audio_path = ov_result[0]

openvino_tokenizer.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58ad5e7e3b08489a7b49897c656adac0183aa72f6ac0f8b146498eb75d4a22e8
+size 4816009

openvino_tokenizer.xml

@@ -0,0 +1,849 @@
+<?xml version="1.0"?>
+<net name="tokenizer" version="11">
+	<layers>
+		<layer id="0" name="Parameter_1" type="Parameter" version="opset1">
+			<data shape="?" element_type="string" />
+			<output>
+				<port id="0" precision="STRING" names="Parameter_1">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="1" name="Constant_7" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="0" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="2" name="StringTensorUnpack_2" type="StringTensorUnpack" version="opset15">
+			<input>
+				<port id="0" precision="STRING">
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="U8">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="3" name="ShapeOf_3" type="ShapeOf" version="opset3">
+			<data output_type="i64" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="I64">
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="4" name="Constant_4" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="0" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="5" name="Constant_5" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="0" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="6" name="Gather_6" type="Gather" version="opset8">
+			<data batch_dims="0" />
+			<input>
+				<port id="0" precision="I64">
+					<dim>1</dim>
+				</port>
+				<port id="1" precision="I64" />
+				<port id="2" precision="I64" />
+			</input>
+			<output>
+				<port id="3" precision="I64" />
+			</output>
+		</layer>
+		<layer id="7" name="Constant_8" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="8" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="8" name="Range_9" type="Range" version="opset4">
+			<data output_type="i32" />
+			<input>
+				<port id="0" precision="I64" />
+				<port id="1" precision="I64" />
+				<port id="2" precision="I64" />
+			</input>
+			<output>
+				<port id="3" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="9" name="Constant_10" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="8" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="10" name="Constant_11" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="8" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="11" name="Add_12" type="Add" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="I64" />
+				<port id="1" precision="I64" />
+			</input>
+			<output>
+				<port id="2" precision="I64" />
+			</output>
+		</layer>
+		<layer id="12" name="Constant_13" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="8" size="8" />
+			<output>
+				<port id="0" precision="I64" />
+			</output>
+		</layer>
+		<layer id="13" name="Range_14" type="Range" version="opset4">
+			<data output_type="i32" />
+			<input>
+				<port id="0" precision="I64" />
+				<port id="1" precision="I64" />
+				<port id="2" precision="I64" />
+			</input>
+			<output>
+				<port id="3" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="14" name="Constant_76" type="Const" version="opset1">
+			<data element_type="u8" shape="5627" offset="16" size="5627" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>5627</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="15" name="SpecialTokensSplit_77" type="SpecialTokensSplit" version="extension">
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="5" precision="U8">
+					<dim>5627</dim>
+				</port>
+			</input>
+			<output>
+				<port id="6" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="10" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="11" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="16" name="Constant_79" type="Const" version="opset1">
+			<data element_type="u8" shape="5" offset="5643" size="5" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>5</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="17" name="Constant_81" type="Const" version="opset1">
+			<data element_type="u8" shape="1" offset="5648" size="1" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="18" name="RegexNormalization_82" type="RegexNormalization" version="extension">
+			<data global_replace="true" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>5</dim>
+				</port>
+				<port id="5" precision="U8">
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="6" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="19" name="Constant_84" type="Const" version="opset1">
+			<data element_type="u8" shape="1" offset="5648" size="1" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="20" name="Constant_86" type="Const" version="opset1">
+			<data element_type="u8" shape="3" offset="5649" size="3" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>3</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="21" name="RegexNormalization_87" type="RegexNormalization" version="extension">
+			<data global_replace="true" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>1</dim>
+				</port>
+				<port id="5" precision="U8">
+					<dim>3</dim>
+				</port>
+			</input>
+			<output>
+				<port id="6" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="22" name="Constant_89" type="Const" version="opset1">
+			<data element_type="u8" shape="11" offset="5652" size="11" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>11</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="23" name="Constant_91" type="Const" version="opset1">
+			<data element_type="u8" shape="5" offset="5663" size="5" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>5</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="24" name="RegexNormalization_92" type="RegexNormalization" version="extension">
+			<data global_replace="true" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>11</dim>
+				</port>
+				<port id="5" precision="U8">
+					<dim>5</dim>
+				</port>
+			</input>
+			<output>
+				<port id="6" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="25" name="Constant_94" type="Const" version="opset1">
+			<data element_type="u8" shape="3" offset="5649" size="3" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>3</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="26" name="RegexSplit_95" type="RegexSplit" version="extension">
+			<data behaviour="mergedwithnext" invert="false" max_splits="-1" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="5" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+				<port id="6" precision="U8">
+					<dim>3</dim>
+				</port>
+			</input>
+			<output>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="10" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="11" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="12" precision="BOOL">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="27" name="Constant_97" type="Const" version="opset1">
+			<data element_type="i32" shape="256300" offset="5668" size="1025200" />
+			<output>
+				<port id="0" precision="I32">
+					<dim>256300</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="28" name="Constant_99" type="Const" version="opset1">
+			<data element_type="i32" shape="256300" offset="1030868" size="1025200" />
+			<output>
+				<port id="0" precision="I32">
+					<dim>256300</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="29" name="Constant_101" type="Const" version="opset1">
+			<data element_type="u8" shape="2247329" offset="2056068" size="2247329" />
+			<output>
+				<port id="0" precision="U8">
+					<dim>2247329</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="30" name="Constant_102_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="256300" offset="4303397" size="512600" />
+			<output>
+				<port id="0" precision="FP16">
+					<dim>256300</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="31" name="Constant_102" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>256300</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>256300</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="32" name="UnigramTokenizer_103" type="UnigramTokenizer" version="extension">
+			<data byte_fallback="false" unk_token_id="3" fuse_unk="true" min_score="inf" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="4" precision="U8">
+					<dim>-1</dim>
+				</port>
+				<port id="5" precision="I32">
+					<dim>256300</dim>
+				</port>
+				<port id="6" precision="I32">
+					<dim>256300</dim>
+				</port>
+				<port id="7" precision="U8">
+					<dim>2247329</dim>
+				</port>
+				<port id="8" precision="FP32">
+					<dim>256300</dim>
+				</port>
+			</input>
+			<output>
+				<port id="9" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="10" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="11" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="33" name="Subtract_104" type="Subtract" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="34" name="Constant_105" type="Const" version="opset1">
+			<data element_type="i32" shape="" offset="4815997" size="4" />
+			<output>
+				<port id="0" precision="I32" />
+			</output>
+		</layer>
+		<layer id="35" name="Minimum_106" type="Minimum" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32" />
+			</input>
+			<output>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="36" name="Add_107" type="Add" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="37" name="Constant_108" type="Const" version="opset1">
+			<data element_type="i32" shape="" offset="4816001" size="4" />
+			<output>
+				<port id="0" precision="I32" />
+			</output>
+		</layer>
+		<layer id="38" name="Constant_109" type="Const" version="opset1">
+			<data element_type="i32" shape="" offset="4816005" size="4" />
+			<output>
+				<port id="0" precision="I32" />
+			</output>
+		</layer>
+		<layer id="39" name="Constant_110" type="Const" version="opset1">
+			<data element_type="i32" shape="1" offset="4816005" size="4" />
+			<output>
+				<port id="0" precision="I32">
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="40" name="Constant_111" type="Const" version="opset1">
+			<data element_type="i32" shape="2" offset="0" size="8" />
+			<output>
+				<port id="0" precision="I32">
+					<dim>2</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="41" name="CombineSegments_112" type="CombineSegments" version="extension">
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="I32" />
+				<port id="4" precision="I32" />
+				<port id="5" precision="I32">
+					<dim>1</dim>
+				</port>
+				<port id="6" precision="I32">
+					<dim>2</dim>
+				</port>
+			</input>
+			<output>
+				<port id="7" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="8" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="9" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="10" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="11" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="12" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="42" name="Subtract_113" type="Subtract" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="43" name="Constant_114" type="Const" version="opset1">
+			<data element_type="i32" shape="" offset="4816001" size="4" />
+			<output>
+				<port id="0" precision="I32" />
+			</output>
+		</layer>
+		<layer id="44" name="ReduceMax_115" type="ReduceMax" version="opset1">
+			<data keep_dims="false" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32" />
+			</input>
+			<output>
+				<port id="2" precision="I32" />
+			</output>
+		</layer>
+		<layer id="45" name="Constant_116" type="Const" version="opset1">
+			<data element_type="i32" shape="" offset="4816001" size="4" />
+			<output>
+				<port id="0" precision="I32" />
+			</output>
+		</layer>
+		<layer id="46" name="RaggedToDense_117" type="RaggedToDense" version="extension">
+			<data pad_right="true" m_pad_max_length="false" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="2" precision="I32">
+					<dim>-1</dim>
+				</port>
+				<port id="3" precision="I32" />
+				<port id="4" precision="I32" />
+			</input>
+			<output>
+				<port id="5" precision="I32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="6" precision="BOOL">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="47" name="Convert_118" type="Convert" version="opset1">
+			<data destination_type="i32" />
+			<input>
+				<port id="0" precision="BOOL">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="I32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="48" name="Convert_118.0" type="Convert" version="opset1">
+			<data destination_type="i64" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="I64" names="attention_mask">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="50" name="RaggedToDense_117.0" type="Convert" version="opset1">
+			<data destination_type="i64" />
+			<input>
+				<port id="0" precision="I32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="I64" names="input_ids">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="51" name="Result_121" type="Result" version="opset1" output_names="input_ids">
+			<input>
+				<port id="0" precision="I64">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+		</layer>
+		<layer id="49" name="Result_123" type="Result" version="opset1" output_names="attention_mask">
+			<input>
+				<port id="0" precision="I64">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+		</layer>
+	</layers>
+	<edges>
+		<edge from-layer="0" from-port="0" to-layer="2" to-port="0" />
+		<edge from-layer="1" from-port="0" to-layer="8" to-port="0" />
+		<edge from-layer="2" from-port="1" to-layer="3" to-port="0" />
+		<edge from-layer="2" from-port="3" to-layer="15" to-port="4" />
+		<edge from-layer="2" from-port="2" to-layer="15" to-port="3" />
+		<edge from-layer="2" from-port="1" to-layer="15" to-port="2" />
+		<edge from-layer="3" from-port="1" to-layer="6" to-port="0" />
+		<edge from-layer="4" from-port="0" to-layer="6" to-port="1" />
+		<edge from-layer="5" from-port="0" to-layer="6" to-port="2" />
+		<edge from-layer="6" from-port="3" to-layer="8" to-port="1" />
+		<edge from-layer="6" from-port="3" to-layer="11" to-port="0" />
+		<edge from-layer="7" from-port="0" to-layer="8" to-port="2" />
+		<edge from-layer="8" from-port="3" to-layer="15" to-port="0" />
+		<edge from-layer="9" from-port="0" to-layer="13" to-port="0" />
+		<edge from-layer="10" from-port="0" to-layer="11" to-port="1" />
+		<edge from-layer="11" from-port="2" to-layer="13" to-port="1" />
+		<edge from-layer="12" from-port="0" to-layer="13" to-port="2" />
+		<edge from-layer="13" from-port="3" to-layer="15" to-port="1" />
+		<edge from-layer="14" from-port="0" to-layer="15" to-port="5" />
+		<edge from-layer="15" from-port="8" to-layer="18" to-port="0" />
+		<edge from-layer="15" from-port="9" to-layer="18" to-port="1" />
+		<edge from-layer="15" from-port="10" to-layer="18" to-port="2" />
+		<edge from-layer="15" from-port="11" to-layer="18" to-port="3" />
+		<edge from-layer="15" from-port="7" to-layer="26" to-port="1" />
+		<edge from-layer="15" from-port="6" to-layer="26" to-port="0" />
+		<edge from-layer="16" from-port="0" to-layer="18" to-port="4" />
+		<edge from-layer="17" from-port="0" to-layer="18" to-port="5" />
+		<edge from-layer="18" from-port="7" to-layer="21" to-port="1" />
+		<edge from-layer="18" from-port="8" to-layer="21" to-port="2" />
+		<edge from-layer="18" from-port="9" to-layer="21" to-port="3" />
+		<edge from-layer="18" from-port="6" to-layer="21" to-port="0" />
+		<edge from-layer="19" from-port="0" to-layer="21" to-port="4" />
+		<edge from-layer="20" from-port="0" to-layer="21" to-port="5" />
+		<edge from-layer="21" from-port="6" to-layer="24" to-port="0" />
+		<edge from-layer="21" from-port="7" to-layer="24" to-port="1" />
+		<edge from-layer="21" from-port="8" to-layer="24" to-port="2" />
+		<edge from-layer="21" from-port="9" to-layer="24" to-port="3" />
+		<edge from-layer="22" from-port="0" to-layer="24" to-port="4" />
+		<edge from-layer="23" from-port="0" to-layer="24" to-port="5" />
+		<edge from-layer="24" from-port="9" to-layer="26" to-port="5" />
+		<edge from-layer="24" from-port="7" to-layer="26" to-port="3" />
+		<edge from-layer="24" from-port="6" to-layer="26" to-port="2" />
+		<edge from-layer="24" from-port="8" to-layer="26" to-port="4" />
+		<edge from-layer="25" from-port="0" to-layer="26" to-port="6" />
+		<edge from-layer="26" from-port="7" to-layer="32" to-port="0" />
+		<edge from-layer="26" from-port="8" to-layer="32" to-port="1" />
+		<edge from-layer="26" from-port="9" to-layer="32" to-port="2" />
+		<edge from-layer="26" from-port="10" to-layer="32" to-port="3" />
+		<edge from-layer="26" from-port="11" to-layer="32" to-port="4" />
+		<edge from-layer="27" from-port="0" to-layer="32" to-port="5" />
+		<edge from-layer="28" from-port="0" to-layer="32" to-port="6" />
+		<edge from-layer="29" from-port="0" to-layer="32" to-port="7" />
+		<edge from-layer="30" from-port="0" to-layer="31" to-port="0" />
+		<edge from-layer="31" from-port="1" to-layer="32" to-port="8" />
+		<edge from-layer="32" from-port="10" to-layer="33" to-port="0" />
+		<edge from-layer="32" from-port="9" to-layer="33" to-port="1" />
+		<edge from-layer="32" from-port="9" to-layer="36" to-port="0" />
+		<edge from-layer="32" from-port="9" to-layer="41" to-port="0" />
+		<edge from-layer="32" from-port="11" to-layer="41" to-port="2" />
+		<edge from-layer="33" from-port="2" to-layer="35" to-port="0" />
+		<edge from-layer="34" from-port="0" to-layer="35" to-port="1" />
+		<edge from-layer="35" from-port="2" to-layer="36" to-port="1" />
+		<edge from-layer="36" from-port="2" to-layer="41" to-port="1" />
+		<edge from-layer="37" from-port="0" to-layer="41" to-port="3" />
+		<edge from-layer="38" from-port="0" to-layer="41" to-port="4" />
+		<edge from-layer="39" from-port="0" to-layer="41" to-port="5" />
+		<edge from-layer="40" from-port="0" to-layer="41" to-port="6" />
+		<edge from-layer="41" from-port="8" to-layer="46" to-port="1" />
+		<edge from-layer="41" from-port="9" to-layer="46" to-port="2" />
+		<edge from-layer="41" from-port="7" to-layer="46" to-port="0" />
+		<edge from-layer="41" from-port="7" to-layer="42" to-port="1" />
+		<edge from-layer="41" from-port="8" to-layer="42" to-port="0" />
+		<edge from-layer="42" from-port="2" to-layer="44" to-port="0" />
+		<edge from-layer="43" from-port="0" to-layer="44" to-port="1" />
+		<edge from-layer="44" from-port="2" to-layer="46" to-port="3" />
+		<edge from-layer="45" from-port="0" to-layer="46" to-port="4" />
+		<edge from-layer="46" from-port="6" to-layer="47" to-port="0" />
+		<edge from-layer="46" from-port="5" to-layer="50" to-port="0" />
+		<edge from-layer="47" from-port="1" to-layer="48" to-port="0" />
+		<edge from-layer="48" from-port="1" to-layer="49" to-port="0" />
+		<edge from-layer="50" from-port="1" to-layer="51" to-port="0" />
+	</edges>
+	<rt_info>
+		<add_attention_mask value="True" />
+		<add_prefix_space />
+		<add_special_tokens value="True" />
+		<bos_token_id value="2" />
+		<clean_up_tokenization_spaces />
+		<detokenizer_input_type value="i64" />
+		<eos_token_id value="1" />
+		<handle_special_tokens_with_re />
+		<max_length />
+		<number_of_inputs value="1" />
+		<openvino_tokenizers_version value="2025.1.0.0-523-710ddf14de8" />
+		<openvino_version value="2025.1.0-18503-6fec06580ab-releases/2025/1" />
+		<original_post_processor_template value="{&quot;type&quot;: &quot;TemplateProcessing&quot;, &quot;single&quot;: [{&quot;Sequence&quot;: {&quot;id&quot;: &quot;A&quot;, &quot;type_id&quot;: 0}}, {&quot;SpecialToken&quot;: {&quot;id&quot;: &quot;&lt;/s>&quot;, &quot;type_id&quot;: 0}}], &quot;pair&quot;: [{&quot;Sequence&quot;: {&quot;id&quot;: &quot;A&quot;, &quot;type_id&quot;: 0}}, {&quot;SpecialToken&quot;: {&quot;id&quot;: &quot;&lt;/s>&quot;, &quot;type_id&quot;: 0}}, {&quot;Sequence&quot;: {&quot;id&quot;: &quot;B&quot;, &quot;type_id&quot;: 0}}, {&quot;SpecialToken&quot;: {&quot;id&quot;: &quot;&lt;/s>&quot;, &quot;type_id&quot;: 0}}], &quot;special_tokens&quot;: {&quot;&lt;/s>&quot;: {&quot;id&quot;: &quot;&lt;/s>&quot;, &quot;ids&quot;: [1], &quot;tokens&quot;: [&quot;&lt;/s>&quot;]}}}" />
+		<original_tokenizer_class value="&lt;class 'transformers.models.t5.tokenization_t5_fast.T5TokenizerFast'>" />
+		<pad_token_id value="0" />
+		<processed_post_processor_template value="{&quot;single&quot;: {&quot;ids&quot;: [-1, 1], &quot;type_ids&quot;: [0, 0]}, &quot;pair&quot;: {&quot;ids&quot;: [-1, 1, -2, 1], &quot;type_ids&quot;: [0, 0, 0, 0]}}" />
+		<sentencepiece_version value="0.2.1" />
+		<skip_special_tokens value="True" />
+		<streaming_detokenizer value="False" />
+		<tiktoken_version value="0.12.0" />
+		<tokenizer_output_type value="i64" />
+		<tokenizers_version value="0.22.1" />
+		<transformers_version value="4.57.3" />
+		<use_max_padding value="False" />
+		<use_sentencepiece_backend value="False" />
+		<utf8_replace_mode value="replace" />
+		<with_detokenizer value="False" />
+	</rt_info>
+</net>

ov_ace_helper.py

@@ -0,0 +1,978 @@
+import os
+import gc
+import math
+import torch
+import types
+import torchaudio
+import torchvision.transforms as transforms
+
+from tqdm import tqdm
+from pathlib import Path
+from loguru import logger
+from diffusers.utils.torch_utils import randn_tensor
+from typing import Dict, Optional, List, Union, Type
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
+
+import nncf
+import openvino as ov
+from openvino.tools.ovc import convert_model
+from openvino_tokenizers import convert_tokenizer
+from openvino.frontend.pytorch.patch_model import __make_16bit_traceable
+
+from acestep.language_segmentation import LangSegment, language_filters
+from acestep.models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
+
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.models.ace_step_transformer import Transformer2DModelOutput
+from acestep.music_dcae.music_dcae_pipeline import MusicDCAE
+from acestep.schedulers.scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler
+from acestep.schedulers.scheduling_flow_match_pingpong import FlowMatchPingPongScheduler
+from acestep.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from acestep.apg_guidance import (
+    apg_forward,
+    MomentumBuffer,
+    cfg_forward,
+    cfg_zero_star,
+    cfg_double_condition_forward,
+)
+
+torch.set_float32_matmul_precision("high")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+TOKENIZER_MODEL_NAME = "openvino_tokenizer.xml"
+TEXT_ENCODER_MODEL_NAME = "ov_text_encoder_model.xml"
+DCAE_ENCODER_MODEL_NAME = "ov_dcae_encoder_model.xml"
+DCAE_DECODER_MODEL_NAME = "ov_dcae_decoder_model.xml"
+VOCODER_DECODE_MODEL_NAME = "ov_vocoder_decode_model.xml"
+VOCODER_MEL_TRANSFORM_MODEL_NAME = "ov_vocoder_mel_transform_model.xml"
+TRANSFORMER_DECODER_MODEL_NAME = "ov_transformer_decoder_model.xml"
+TRANSFORMER_ENCODER_MODEL_NAME = "ov_transformer_encoder_model.xml"
+
+
+def cleanup_torchscript_cache():
+    """
+    Helper for removing cached model representation
+    """
+    torch._C._jit_clear_class_registry()
+    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
+    torch.jit._state._clear_class_state()
+
+
+def ov_convert(
+    model_dir_path: str,
+    ov_model_name: str,
+    inputs: Dict,
+    orig_model: Type[torch.nn.Module],
+    model_name: str,
+    quantization_config: Dict = None,
+    force_convertion: bool = False,
+):
+    try:
+        ov_model_path = Path(model_dir_path, ov_model_name)
+        if not ov_model_path.exists() or force_convertion:
+            print(f"⌛ Convert {model_name} model")
+            orig_model.eval()
+            __make_16bit_traceable(orig_model)
+            ov_model = convert_model(orig_model, example_input=inputs)
+            if quantization_config is not None:
+                print(f"⌛ Weights compression with {quantization_config['mode']} mode started")
+                ov_model = nncf.compress_weights(ov_model, **quantization_config)
+                print("✅ Weights compression finished")
+            ov.save_model(ov_model, ov_model_path)
+
+            del ov_model
+            cleanup_torchscript_cache()
+            gc.collect()
+            print(f"✅ {model_name} model converted")
+    except Exception as e:
+        print(f"❌{model_name} model is not converted. Error: {e}")
+
+
+def convert_transformer_models(pipeline: ACEStepPipeline, model_dir: str = "ov_converted", orig_checkpoint_path: str = "", quantization_config: Dict = None):
+    # Transformer Encoder model
+    def encode_with_temperature_wrap(
+        self,
+        encoder_text_hidden_states: torch.Tensor = None,
+        text_attention_mask: torch.LongTensor = None,
+        speaker_embeds: torch.FloatTensor = None,
+        lyric_token_idx: torch.LongTensor = None,
+        lyric_mask: torch.LongTensor = None,
+        tau: torch.FloatTensor = torch.Tensor([0.01]),
+    ):
+        handlers = []
+
+        def hook(module, input, output):
+            output[:] *= tau[0]
+            return output
+
+        l_min = 4
+        l_max = 6
+        for i in range(l_min, l_max):
+            handler = self.lyric_encoder.encoders[i].self_attn.linear_q.register_forward_hook(hook)
+            handlers.append(handler)
+
+        encoder_hidden_states, encoder_hidden_mask = self.encode(
+            encoder_text_hidden_states=encoder_text_hidden_states,
+            text_attention_mask=text_attention_mask,
+            speaker_embeds=speaker_embeds,
+            lyric_token_idx=lyric_token_idx,
+            lyric_mask=lyric_mask,
+        )
+
+        for hook in handlers:
+            hook.remove()
+
+        return encoder_hidden_states, encoder_hidden_mask
+
+    inputs = {
+        "encoder_text_hidden_states": torch.randn(size=(1, 15, 768), dtype=torch.float),
+        "text_attention_mask": torch.ones([1, 15], dtype=torch.int64),
+        "speaker_embeds": torch.zeros(size=(1, 512), dtype=torch.float),
+        "lyric_token_idx": torch.randint(10000, [1, 543], dtype=torch.int64),
+        "lyric_mask": torch.ones([1, 543], dtype=torch.int64),
+        "tau": torch.Tensor([0.01]),
+    }
+    transformer_encoder_model = pipeline.ace_step_transformer
+    transformer_encoder_erg_model = pipeline.ace_step_transformer
+    transformer_encoder_erg_model.forward = types.MethodType(encode_with_temperature_wrap, transformer_encoder_model)
+    ov_convert(
+        model_dir,
+        TRANSFORMER_ENCODER_MODEL_NAME,
+        inputs,
+        transformer_encoder_erg_model,
+        "Transformer Encoder with Entropy Rectifying Guidance",
+        quantization_config=quantization_config,
+    )
+
+    # Transformer Decoder model
+    def decode_with_temperature_wrap(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_mask: torch.Tensor,
+        timestep: torch.Tensor = None,
+        # ssl_hidden_states: List[torch.Tensor] = None,
+        output_length: int = 0,
+        # block_controlnet_hidden_states: Union[List[torch.Tensor], torch.Tensor] = None,
+        # controlnet_scale: Union[float, torch.Tensor] = 1.0,
+        tau: torch.FloatTensor = torch.Tensor([0.01]),
+    ):
+        handlers = []
+
+        def hook(module, input, output):
+            output[:] *= tau[0]
+            return output
+
+        l_min = 5
+        l_max = 10
+        for i in range(l_min, l_max):
+            handler = self.transformer_blocks[i].attn.to_q.register_forward_hook(hook)
+            handlers.append(handler)
+            handler = self.transformer_blocks[i].cross_attn.to_q.register_forward_hook(hook)
+            handlers.append(handler)
+
+        sample = self.decode(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_hidden_mask=encoder_hidden_mask,
+            output_length=output_length,
+            timestep=timestep,
+        ).sample
+
+        for hook in handlers:
+            hook.remove()
+
+        return sample
+
+    inputs = {
+        "hidden_states": torch.randn(size=(1, 8, 16, 151), dtype=torch.float),
+        "attention_mask": torch.ones([1, 151], dtype=torch.int64),
+        "encoder_hidden_states": torch.randn(size=(1, 559, 2560), dtype=torch.float),
+        "encoder_hidden_mask": torch.ones([1, 559], dtype=torch.float),
+        "output_length": torch.tensor(151),
+        "timestep": torch.randn([1], dtype=torch.float),
+        "tau": torch.Tensor([0.01]),
+    }
+    transformer_decoder_erg_model = pipeline.ace_step_transformer
+    transformer_decoder_erg_model.forward = types.MethodType(decode_with_temperature_wrap, transformer_decoder_erg_model)
+    ov_convert(
+        model_dir,
+        TRANSFORMER_DECODER_MODEL_NAME,
+        inputs,
+        transformer_decoder_erg_model,
+        "Transformer Decoder with Entropy Rectifying Guidance",
+        quantization_config=quantization_config,
+    )
+
+
+def convert_models(pipeline: ACEStepPipeline, model_dir: str = "ov_converted_new", orig_checkpoint_path: str = "", quantization_config: Dict = None):
+    print(f"⌛ Conversion started. Be patient, it may takes some time.")
+
+    if not pipeline.loaded or (orig_checkpoint_path and not Path(orig_checkpoint_path).exists()):
+        print("⌛ Load Original model checkpoints")
+        pipeline.load_checkpoint(orig_checkpoint_path)
+        print("✅ Original model checkpoints successfully loaded")
+
+    # Tokenizer
+    ov_tokenizer_path = Path(model_dir, TOKENIZER_MODEL_NAME)
+    if not ov_tokenizer_path.exists():
+        print(f"⌛ Convert Tokenizer")
+        if not ov_tokenizer_path.exists():
+            ov_tokenizer = convert_tokenizer(pipeline.text_tokenizer, with_detokenizer=False)
+            ov.save_model(ov_tokenizer, Path(model_dir, TOKENIZER_MODEL_NAME))
+        print(f"✅ Tokenizer is converted")
+
+    # Text Encoder Model
+    inputs = {
+        "input_ids": torch.randint(1000, size=(1, 15), dtype=torch.int64),
+        "attention_mask": torch.ones([1, 15], dtype=torch.int64),
+    }
+    ov_convert(model_dir, TEXT_ENCODER_MODEL_NAME, inputs, pipeline.text_encoder_model, "UMT5 Encoder")
+
+    # DCAE Encoder model
+    inputs = {"hidden_states": torch.randn([1, 2, 128, 1208], dtype=torch.float)}
+    ov_convert(model_dir, DCAE_ENCODER_MODEL_NAME, inputs, pipeline.music_dcae.dcae.encoder, "Sana's Deep Compression AutoEncoder")
+
+    # DCAE Decoder model
+    inputs = {"hidden_states": torch.randn([1, 8, 16, 151], dtype=torch.float)}
+    ov_convert(model_dir, DCAE_DECODER_MODEL_NAME, inputs, pipeline.music_dcae.dcae.decoder, "Sana's Deep Compression AutoEncoder Decoder")
+
+    # Vocoder Mel Transform model
+    inputs = {"x": torch.randn([2, 618496], dtype=torch.float)}
+    ov_convert(model_dir, VOCODER_MEL_TRANSFORM_MODEL_NAME, inputs, pipeline.music_dcae.vocoder.mel_transform, "Vocoder Mel Transform")
+
+    # Vocoder Decoder model
+    inputs = {"mel": torch.randn([1, 128, 856], dtype=torch.float)}
+    ov_convert(model_dir, VOCODER_DECODE_MODEL_NAME, inputs, pipeline.music_dcae.vocoder, "Vocoder Decoder")
+
+    # DiT
+    convert_transformer_models(pipeline, model_dir, orig_checkpoint_path, quantization_config)
+
+
+class MusicDCAEWrapper(MusicDCAE):
+    def __init__(self, source_sample_rate=None):
+        torch.nn.Module.__init__(self)
+        self.dcae = None
+        self.vocoder = None
+
+        if source_sample_rate is None:
+            source_sample_rate = 48000
+
+        self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Normalize(0.5, 0.5),
+            ]
+        )
+        self.min_mel_value = -11.0
+        self.max_mel_value = 3.0
+        self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
+        self.mel_chunk_size = 1024
+        self.time_dimention_multiple = 8
+        self.latent_chunk_size = self.mel_chunk_size // self.time_dimention_multiple
+        self.scale_factor = 0.1786
+        self.shift_factor = -1.9091
+
+
+class OVDCAECompiledModels(torch.nn.Module):
+    def __init__(self, compiled_model):
+        self.compiled_model = compiled_model
+
+    def __call__(self, inputs):
+        if not self.compiled_model:
+            logger.error("OVDCAECompiledModels: compiled model is not defined")
+
+        output = self.compiled_model({"hidden_states": inputs.to(dtype=torch.float32)})
+        return torch.from_numpy(output[0])
+
+    @classmethod
+    def from_pretrained(cls, ov_model_path, device, ov_core):
+        ov_dcae_model = ov_core.read_model(ov_model_path)
+        compiled_model = ov_core.compile_model(ov_dcae_model, device)
+        return cls(compiled_model)
+
+
+class OVWrapperAutoencoderDC(torch.nn.Module):
+    def __init__(self, encoder, decoder):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+
+    @classmethod
+    def from_pretrained(cls, ov_core, ov_models_path, device="CPU"):
+        encoder = OVDCAECompiledModels.from_pretrained(Path(ov_models_path, DCAE_ENCODER_MODEL_NAME), device, ov_core)
+        decoder = OVDCAECompiledModels.from_pretrained(Path(ov_models_path, DCAE_DECODER_MODEL_NAME), device, ov_core)
+        return cls(encoder, decoder)
+
+
+class OVWrapperADaMoSHiFiGANV1(torch.nn.Module):
+    def __init__(self, encoder_compiled_model, mel_trnasform_compiled_model):
+        super().__init__()
+        self.decoder = encoder_compiled_model
+        self.mel_trnasform = mel_trnasform_compiled_model
+
+    @classmethod
+    def from_pretrained(cls, ov_core, ov_models_path, device="CPU"):
+        ov_vocoder_decoder_model = ov_core.read_model(Path(ov_models_path, VOCODER_DECODE_MODEL_NAME))
+        decoder = ov_core.compile_model(ov_vocoder_decoder_model, device)
+        ov_vocoder_mel_transform_model = ov_core.read_model(Path(ov_models_path, VOCODER_MEL_TRANSFORM_MODEL_NAME))
+        mel_trnasform = ov_core.compile_model(ov_vocoder_mel_transform_model, device)
+        return cls(decoder, mel_trnasform)
+
+    def decode(self, inputs):
+        output = self.decoder({"mel": inputs.to(dtype=torch.float32)})
+        return torch.from_numpy(output[0])
+
+    def mel_transform(self, inputs):
+        output = self.mel_trnasform({"x": inputs.to(dtype=torch.float32)})
+        return torch.from_numpy(output[0])
+
+    def forward(self, inputs):
+        return self.decode(inputs)
+
+
+class OvWrapperACEStepTransformer2DModel(torch.nn.Module):
+    def __init__(self, encoder_model, decoder_model):
+        super().__init__()
+        self.ov_lyric_encoder_compiled = encoder_model
+        self.ov_decoder_compiled_model = decoder_model
+
+    @classmethod
+    def from_pretrained(cls, ov_core, ov_models_path, device="CPU"):
+        ov_model_encoder = ov_core.read_model(Path(ov_models_path, TRANSFORMER_ENCODER_MODEL_NAME))
+        compiled_model_encoder = ov_core.compile_model(ov_model_encoder, device)
+
+        ov_model_decoder = ov_core.read_model(Path(ov_models_path, TRANSFORMER_DECODER_MODEL_NAME))
+        compiled_model_decoder = ov_core.compile_model(ov_model_decoder, device)
+        return cls(compiled_model_encoder, compiled_model_decoder)
+
+    def encode_with_temperature(
+        self,
+        encoder_text_hidden_states: Optional[torch.Tensor] = None,
+        text_attention_mask: Optional[torch.LongTensor] = None,
+        speaker_embeds: Optional[torch.FloatTensor] = None,
+        lyric_token_idx: Optional[torch.LongTensor] = None,
+        lyric_mask: Optional[torch.LongTensor] = None,
+        tau: Optional[torch.FloatTensor] = torch.Tensor([0.01]),
+    ):
+        output = None
+        if self.ov_lyric_encoder_compiled:
+            output = self.ov_lyric_encoder_compiled(
+                {
+                    "encoder_text_hidden_states": encoder_text_hidden_states,
+                    "text_attention_mask": text_attention_mask,
+                    "speaker_embeds": speaker_embeds,
+                    "lyric_token_idx": lyric_token_idx,
+                    "lyric_mask": lyric_mask,
+                    "tau": tau,
+                }
+            )
+        return torch.from_numpy(output[0]), torch.from_numpy(output[1])
+
+    def decode_with_temperature(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_mask: torch.Tensor,
+        timestep: Optional[torch.Tensor],
+        ssl_hidden_states: Optional[List[torch.Tensor]] = None,
+        output_length: int = 0,
+        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        controlnet_scale: Union[float, torch.Tensor] = 1.0,
+        return_dict: bool = True,
+        tau: Optional[torch.FloatTensor] = torch.Tensor([0.01]),
+    ):
+        output = None
+        if self.ov_decoder_compiled_model:
+            output = self.ov_decoder_compiled_model(
+                {
+                    "hidden_states": hidden_states,
+                    "attention_mask": attention_mask,
+                    "encoder_hidden_states": encoder_hidden_states,
+                    "encoder_hidden_mask": encoder_hidden_mask,
+                    "output_length": output_length,
+                    "timestep": timestep,
+                    "tau": tau,
+                }
+            )
+
+        sample = torch.from_numpy(output[0]) if output is not None else None
+        return sample
+
+    def encode(
+        self,
+        encoder_text_hidden_states: Optional[torch.Tensor] = None,
+        text_attention_mask: Optional[torch.LongTensor] = None,
+        speaker_embeds: Optional[torch.FloatTensor] = None,
+        lyric_token_idx: Optional[torch.LongTensor] = None,
+        lyric_mask: Optional[torch.LongTensor] = None,
+    ):
+        return self.encode_with_temperature(
+            encoder_text_hidden_states=encoder_text_hidden_states,
+            text_attention_mask=text_attention_mask,
+            speaker_embeds=speaker_embeds,
+            lyric_token_idx=lyric_token_idx,
+            lyric_mask=lyric_mask,
+            tau=torch.Tensor([1]),
+        )
+
+    def decode(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_mask: torch.Tensor,
+        timestep: Optional[torch.Tensor],
+        ssl_hidden_states: Optional[List[torch.Tensor]] = None,
+        output_length: int = 0,
+        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        controlnet_scale: Union[float, torch.Tensor] = 1.0,
+        return_dict: bool = True,
+    ):
+        sample = self.decode_with_temperature(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_hidden_mask=encoder_hidden_mask,
+            timestep=timestep,
+            ssl_hidden_states=ssl_hidden_states,
+            output_length=output_length,
+            block_controlnet_hidden_states=block_controlnet_hidden_states,
+            controlnet_scale=controlnet_scale,
+            return_dict=return_dict,
+            tau=torch.Tensor([1]),
+        )
+
+        return Transformer2DModelOutput(sample, None)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_text_hidden_states: Optional[torch.Tensor] = None,
+        text_attention_mask: Optional[torch.LongTensor] = None,
+        speaker_embeds: Optional[torch.FloatTensor] = None,
+        lyric_token_idx: Optional[torch.LongTensor] = None,
+        lyric_mask: Optional[torch.LongTensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+        ssl_hidden_states: Optional[List[torch.Tensor]] = None,
+        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        controlnet_scale: Union[float, torch.Tensor] = 1.0,
+        return_dict: bool = True,
+    ):
+        encoder_hidden_states, encoder_hidden_mask = self.encode(
+            encoder_text_hidden_states=encoder_text_hidden_states,
+            text_attention_mask=text_attention_mask,
+            speaker_embeds=speaker_embeds,
+            lyric_token_idx=lyric_token_idx,
+            lyric_mask=lyric_mask,
+        )
+
+        output_length = hidden_states.shape[-1]
+
+        output = self.decode(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_hidden_mask=encoder_hidden_mask,
+            timestep=timestep,
+            ssl_hidden_states=ssl_hidden_states,
+            output_length=output_length,
+            block_controlnet_hidden_states=block_controlnet_hidden_states,
+            controlnet_scale=controlnet_scale,
+            return_dict=return_dict,
+        )
+
+        return output
+
+
+class OVACEStepPipeline(ACEStepPipeline):
+    def __init__(self):
+        super().__init__(checkpoint_dir="", dtype="float32")
+        self.core = ov.Core()
+
+        self.dcae_decoder = None
+        self.vocoder_encode = None
+        self.vocoder_decoder = None
+        self.transformer_encode = None
+        self.transformer_encode_with_temperature = None
+        self.transformer_decode = None
+        self.transformer_decode_with_temperature = None
+
+        self.ace_step_transformer_origin = None
+        self.ace_step_transformer = None
+        self.music_dcae = None
+        self.text_tokenizer = None
+        self.text_encoder_model = None
+
+    def get_checkpoint_path(self, checkpoint_dir, repo):
+        pass
+
+    def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False):
+        pass
+
+    def load_models(self, ov_models_path: str = None, device: str = "CPU"):
+        self.loaded = True
+        if ov_models_path and Path(ov_models_path).exists:
+            ov_text_encoder_model = self.core.read_model(Path(ov_models_path, TEXT_ENCODER_MODEL_NAME))
+            self.text_encoder_model = self.core.compile_model(ov_text_encoder_model, device)
+
+            ov_text_tokenizer_path = self.core.read_model(Path(ov_models_path, TOKENIZER_MODEL_NAME))
+            self.text_tokenizer = self.core.compile_model(ov_text_tokenizer_path, "CPU")  # tokenizer can only be inferred on CPU
+
+            self.music_dcae = MusicDCAEWrapper()
+            self.music_dcae.dcae = OVWrapperAutoencoderDC.from_pretrained(self.core, ov_models_path, device)
+            self.music_dcae.vocoder = OVWrapperADaMoSHiFiGANV1.from_pretrained(self.core, ov_models_path, device)
+
+            self.ace_step_transformer = OvWrapperACEStepTransformer2DModel.from_pretrained(self.core, ov_models_path, device)
+        else:
+            logger.error(f"Path is not exists: {ov_models_path}")
+
+        lang_segment = LangSegment()
+        lang_segment.setfilters(language_filters.default)
+        self.lang_segment = lang_segment
+        self.lyric_tokenizer = VoiceBpeTokenizer()
+
+    def load_quantized_checkpoint(self, checkpoint_dir=None):
+        pass
+
+    def get_text_embeddings(self, texts, text_max_length=256):
+        inputs = self.text_tokenizer(texts)
+        inputs = {"attention_mask": inputs["attention_mask"], "input_ids": inputs["input_ids"]}
+
+        last_hidden_states = self.text_encoder_model(inputs)
+        attention_mask = inputs["attention_mask"]
+        return torch.from_numpy(last_hidden_states[0]), torch.from_numpy(attention_mask)
+
+    def get_text_embeddings_null(self, texts, text_max_length=256, tau=0.01, l_min=8, l_max=10):
+        inputs = self.text_tokenizer(texts)
+        inputs = {"attention_mask": inputs["attention_mask"], "input_ids": inputs["input_ids"]}
+        last_hidden_states = self.text_encoder_model(inputs)
+        return torch.from_numpy(last_hidden_states[0])
+
+    def text2music_diffusion_process(
+        self,
+        duration,
+        encoder_text_hidden_states,
+        text_attention_mask,
+        speaker_embds,
+        lyric_token_ids,
+        lyric_mask,
+        random_generators=None,
+        infer_steps=60,
+        guidance_scale=15.0,
+        omega_scale=10.0,
+        scheduler_type="euler",
+        cfg_type="apg",
+        zero_steps=1,
+        use_zero_init=True,
+        guidance_interval=0.5,
+        guidance_interval_decay=1.0,
+        min_guidance_scale=3.0,
+        oss_steps=[],
+        encoder_text_hidden_states_null=None,
+        use_erg_lyric=False,
+        use_erg_diffusion=False,
+        retake_random_generators=None,
+        retake_variance=0.5,
+        add_retake_noise=False,
+        guidance_scale_text=0.0,
+        guidance_scale_lyric=0.0,
+        repaint_start=0,
+        repaint_end=0,
+        src_latents=None,
+        audio2audio_enable=False,
+        ref_audio_strength=0.5,
+        ref_latents=None,
+    ):
+        logger.info("cfg_type: {}, guidance_scale: {}, omega_scale: {}".format(cfg_type, guidance_scale, omega_scale))
+        do_classifier_free_guidance = True
+        if guidance_scale == 0.0 or guidance_scale == 1.0:
+            do_classifier_free_guidance = False
+
+        do_double_condition_guidance = False
+        if guidance_scale_text is not None and guidance_scale_text > 1.0 and guidance_scale_lyric is not None and guidance_scale_lyric > 1.0:
+            do_double_condition_guidance = True
+            logger.info(
+                "do_double_condition_guidance: {}, guidance_scale_text: {}, guidance_scale_lyric: {}".format(
+                    do_double_condition_guidance,
+                    guidance_scale_text,
+                    guidance_scale_lyric,
+                )
+            )
+
+        bsz = encoder_text_hidden_states.shape[0]
+
+        if scheduler_type == "euler":
+            scheduler = FlowMatchEulerDiscreteScheduler(
+                num_train_timesteps=1000,
+                shift=3.0,
+            )
+        elif scheduler_type == "heun":
+            scheduler = FlowMatchHeunDiscreteScheduler(
+                num_train_timesteps=1000,
+                shift=3.0,
+            )
+        elif scheduler_type == "pingpong":
+            scheduler = FlowMatchPingPongScheduler(
+                num_train_timesteps=1000,
+                shift=3.0,
+            )
+
+        frame_length = int(duration * 44100 / 512 / 8)
+        if src_latents is not None:
+            frame_length = src_latents.shape[-1]
+
+        if ref_latents is not None:
+            frame_length = ref_latents.shape[-1]
+
+        if len(oss_steps) > 0:
+            infer_steps = max(oss_steps)
+            scheduler.set_timesteps
+            timesteps, num_inference_steps = retrieve_timesteps(
+                scheduler,
+                num_inference_steps=infer_steps,
+                device=self.device,
+                timesteps=None,
+            )
+            new_timesteps = torch.zeros(len(oss_steps), dtype=self.dtype, device=self.device)
+            for idx in range(len(oss_steps)):
+                new_timesteps[idx] = timesteps[oss_steps[idx] - 1]
+            num_inference_steps = len(oss_steps)
+            sigmas = (new_timesteps / 1000).float().cpu().numpy()
+            timesteps, num_inference_steps = retrieve_timesteps(
+                scheduler,
+                num_inference_steps=num_inference_steps,
+                device=self.device,
+                sigmas=sigmas,
+            )
+            logger.info(f"oss_steps: {oss_steps}, num_inference_steps: {num_inference_steps} after remapping to timesteps {timesteps}")
+        else:
+            timesteps, num_inference_steps = retrieve_timesteps(
+                scheduler,
+                num_inference_steps=infer_steps,
+                device=self.device,
+                timesteps=None,
+            )
+
+        target_latents = randn_tensor(
+            shape=(bsz, 8, 16, frame_length),
+            generator=random_generators,
+            device=self.device,
+            dtype=self.dtype,
+        )
+
+        is_repaint = False
+        is_extend = False
+
+        if add_retake_noise:
+            n_min = int(infer_steps * (1 - retake_variance))
+            retake_variance = torch.tensor(retake_variance * math.pi / 2).to(self.device).to(self.dtype)
+            retake_latents = randn_tensor(
+                shape=(bsz, 8, 16, frame_length),
+                generator=retake_random_generators,
+                device=self.device,
+                dtype=self.dtype,
+            )
+            repaint_start_frame = int(repaint_start * 44100 / 512 / 8)
+            repaint_end_frame = int(repaint_end * 44100 / 512 / 8)
+            x0 = src_latents
+            # retake
+            is_repaint = repaint_end_frame - repaint_start_frame != frame_length
+
+            is_extend = (repaint_start_frame < 0) or (repaint_end_frame > frame_length)
+            if is_extend:
+                is_repaint = True
+
+            # TODO: train a mask aware repainting controlnet
+            # to make sure mean = 0, std = 1
+            if not is_repaint:
+                target_latents = torch.cos(retake_variance) * target_latents + torch.sin(retake_variance) * retake_latents
+            elif not is_extend:
+                # if repaint_end_frame
+                repaint_mask = torch.zeros((bsz, 8, 16, frame_length), device=self.device, dtype=self.dtype)
+                repaint_mask[:, :, :, repaint_start_frame:repaint_end_frame] = 1.0
+                repaint_noise = torch.cos(retake_variance) * target_latents + torch.sin(retake_variance) * retake_latents
+                repaint_noise = torch.where(repaint_mask == 1.0, repaint_noise, target_latents)
+                zt_edit = x0.clone()
+                z0 = repaint_noise
+            elif is_extend:
+                to_right_pad_gt_latents = None
+                to_left_pad_gt_latents = None
+                gt_latents = src_latents
+                src_latents_length = gt_latents.shape[-1]
+                max_infer_fame_length = int(240 * 44100 / 512 / 8)
+                left_pad_frame_length = 0
+                right_pad_frame_length = 0
+                right_trim_length = 0
+                left_trim_length = 0
+                if repaint_start_frame < 0:
+                    left_pad_frame_length = abs(repaint_start_frame)
+                    frame_length = left_pad_frame_length + gt_latents.shape[-1]
+                    extend_gt_latents = torch.nn.functional.pad(gt_latents, (left_pad_frame_length, 0), "constant", 0)
+                    if frame_length > max_infer_fame_length:
+                        right_trim_length = frame_length - max_infer_fame_length
+                        extend_gt_latents = extend_gt_latents[:, :, :, :max_infer_fame_length]
+                        to_right_pad_gt_latents = extend_gt_latents[:, :, :, -right_trim_length:]
+                        frame_length = max_infer_fame_length
+                    repaint_start_frame = 0
+                    gt_latents = extend_gt_latents
+
+                if repaint_end_frame > src_latents_length:
+                    right_pad_frame_length = repaint_end_frame - gt_latents.shape[-1]
+                    frame_length = gt_latents.shape[-1] + right_pad_frame_length
+                    extend_gt_latents = torch.nn.functional.pad(gt_latents, (0, right_pad_frame_length), "constant", 0)
+                    if frame_length > max_infer_fame_length:
+                        left_trim_length = frame_length - max_infer_fame_length
+                        extend_gt_latents = extend_gt_latents[:, :, :, -max_infer_fame_length:]
+                        to_left_pad_gt_latents = extend_gt_latents[:, :, :, :left_trim_length]
+                        frame_length = max_infer_fame_length
+                    repaint_end_frame = frame_length
+                    gt_latents = extend_gt_latents
+
+                repaint_mask = torch.zeros((bsz, 8, 16, frame_length), device=self.device, dtype=self.dtype)
+                if left_pad_frame_length > 0:
+                    repaint_mask[:, :, :, :left_pad_frame_length] = 1.0
+                if right_pad_frame_length > 0:
+                    repaint_mask[:, :, :, -right_pad_frame_length:] = 1.0
+                x0 = gt_latents
+                padd_list = []
+                if left_pad_frame_length > 0:
+                    padd_list.append(retake_latents[:, :, :, :left_pad_frame_length])
+                padd_list.append(
+                    target_latents[
+                        :,
+                        :,
+                        :,
+                        left_trim_length : target_latents.shape[-1] - right_trim_length,
+                    ]
+                )
+                if right_pad_frame_length > 0:
+                    padd_list.append(retake_latents[:, :, :, -right_pad_frame_length:])
+                target_latents = torch.cat(padd_list, dim=-1)
+                assert target_latents.shape[-1] == x0.shape[-1], f"{target_latents.shape=} {x0.shape=}"
+                zt_edit = x0.clone()
+                z0 = target_latents
+
+        if audio2audio_enable and ref_latents is not None:
+            logger.info(f"audio2audio_enable: {audio2audio_enable}, ref_latents: {ref_latents.shape}")
+            target_latents, timesteps, scheduler, num_inference_steps = self.add_latents_noise(
+                gt_latents=ref_latents,
+                sigma_max=(1 - ref_audio_strength),
+                noise=target_latents,
+                scheduler_type=scheduler_type,
+                infer_steps=infer_steps,
+            )
+
+        attention_mask = torch.ones(bsz, frame_length, device=self.device, dtype=self.dtype)
+
+        # guidance interval
+        start_idx = int(num_inference_steps * ((1 - guidance_interval) / 2))
+        end_idx = int(num_inference_steps * (guidance_interval / 2 + 0.5))
+        logger.info(f"start_idx: {start_idx}, end_idx: {end_idx}, num_inference_steps: {num_inference_steps}")
+
+        momentum_buffer = MomentumBuffer()
+
+        # P(speaker, text, lyric)
+        encoder_hidden_states, encoder_hidden_mask = self.ace_step_transformer.encode(
+            encoder_text_hidden_states,
+            text_attention_mask,
+            speaker_embds,
+            lyric_token_ids,
+            lyric_mask,
+        )
+
+        if use_erg_lyric:
+            # P(null_speaker, text_weaker, lyric_weaker)
+            encoder_hidden_states_null, _ = self.ace_step_transformer.encode_with_temperature(
+                encoder_text_hidden_states=(
+                    encoder_text_hidden_states_null if encoder_text_hidden_states_null is not None else torch.zeros_like(encoder_text_hidden_states)
+                ),
+                text_attention_mask=text_attention_mask,
+                speaker_embeds=torch.zeros_like(speaker_embds),
+                lyric_token_idx=lyric_token_ids,
+                lyric_mask=lyric_mask,
+            )
+        else:
+            # P(null_speaker, null_text, null_lyric)
+            encoder_hidden_states_null, _ = self.ace_step_transformer.encode(
+                torch.zeros_like(encoder_text_hidden_states),
+                text_attention_mask,
+                torch.zeros_like(speaker_embds),
+                torch.zeros_like(lyric_token_ids),
+                lyric_mask,
+            )
+
+        encoder_hidden_states_no_lyric = None
+        if do_double_condition_guidance:
+            # P(null_speaker, text, lyric_weaker)
+            if use_erg_lyric:
+                encoder_hidden_states_no_lyric, _ = self.ace_step_transformer.encode_with_temperature(
+                    encoder_text_hidden_states=encoder_text_hidden_states,
+                    text_attention_mask=text_attention_mask,
+                    speaker_embeds=torch.zeros_like(speaker_embds),
+                    lyric_token_idx=lyric_token_ids,
+                    lyric_mask=lyric_mask,
+                )
+            # P(null_speaker, text, no_lyric)
+            else:
+                encoder_hidden_states_no_lyric, _ = self.ace_step_transformer.encode(
+                    encoder_text_hidden_states,
+                    text_attention_mask,
+                    torch.zeros_like(speaker_embds),
+                    torch.zeros_like(lyric_token_ids),
+                    lyric_mask,
+                )
+
+        for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
+            if is_repaint:
+                if i < n_min:
+                    continue
+                elif i == n_min:
+                    t_i = t / 1000
+                    zt_src = (1 - t_i) * x0 + (t_i) * z0
+                    target_latents = zt_edit + zt_src - x0
+                    logger.info(f"repaint start from {n_min} add {t_i} level of noise")
+
+            # expand the latents if we are doing classifier free guidance
+            latents = target_latents
+
+            is_in_guidance_interval = start_idx <= i < end_idx
+            if is_in_guidance_interval and do_classifier_free_guidance:
+                # compute current guidance scale
+                if guidance_interval_decay > 0:
+                    # Linearly interpolate to calculate the current guidance scale
+                    progress = (i - start_idx) / (end_idx - start_idx - 1)  # 归一化到[0,1]
+                    current_guidance_scale = guidance_scale - (guidance_scale - min_guidance_scale) * progress * guidance_interval_decay
+                else:
+                    current_guidance_scale = guidance_scale
+
+                latent_model_input = latents
+                timestep = t.expand(latent_model_input.shape[0])
+                output_length = latent_model_input.shape[-1]
+                # P(x|speaker, text, lyric)
+                noise_pred_with_cond = self.ace_step_transformer.decode(
+                    hidden_states=latent_model_input,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_mask=encoder_hidden_mask,
+                    output_length=output_length,
+                    timestep=timestep,
+                ).sample
+
+                noise_pred_with_only_text_cond = None
+                if do_double_condition_guidance and encoder_hidden_states_no_lyric is not None:
+                    noise_pred_with_only_text_cond = self.ace_step_transformer.decode(
+                        hidden_states=latent_model_input,
+                        attention_mask=attention_mask,
+                        encoder_hidden_states=encoder_hidden_states_no_lyric,
+                        encoder_hidden_mask=encoder_hidden_mask,
+                        output_length=output_length,
+                        timestep=timestep,
+                    ).sample
+
+                if use_erg_diffusion:
+                    noise_pred_uncond = self.ace_step_transformer.decode_with_temperature(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=encoder_hidden_states_null,
+                        encoder_hidden_mask=encoder_hidden_mask,
+                        output_length=output_length,
+                        attention_mask=attention_mask,
+                    )
+                else:
+                    noise_pred_uncond = self.ace_step_transformer.decode(
+                        hidden_states=latent_model_input,
+                        attention_mask=attention_mask,
+                        encoder_hidden_states=encoder_hidden_states_null,
+                        encoder_hidden_mask=encoder_hidden_mask,
+                        output_length=output_length,
+                        timestep=timestep,
+                    ).sample
+
+                if do_double_condition_guidance and noise_pred_with_only_text_cond is not None:
+                    noise_pred = cfg_double_condition_forward(
+                        cond_output=noise_pred_with_cond,
+                        uncond_output=noise_pred_uncond,
+                        only_text_cond_output=noise_pred_with_only_text_cond,
+                        guidance_scale_text=guidance_scale_text,
+                        guidance_scale_lyric=guidance_scale_lyric,
+                    )
+
+                elif cfg_type == "apg":
+                    noise_pred = apg_forward(
+                        pred_cond=noise_pred_with_cond,
+                        pred_uncond=noise_pred_uncond,
+                        guidance_scale=current_guidance_scale,
+                        momentum_buffer=momentum_buffer,
+                    )
+                elif cfg_type == "cfg":
+                    noise_pred = cfg_forward(
+                        cond_output=noise_pred_with_cond,
+                        uncond_output=noise_pred_uncond,
+                        cfg_strength=current_guidance_scale,
+                    )
+                elif cfg_type == "cfg_star":
+                    noise_pred = cfg_zero_star(
+                        noise_pred_with_cond=noise_pred_with_cond,
+                        noise_pred_uncond=noise_pred_uncond,
+                        guidance_scale=current_guidance_scale,
+                        i=i,
+                        zero_steps=zero_steps,
+                        use_zero_init=use_zero_init,
+                    )
+            else:
+                latent_model_input = latents
+                timestep = t.expand(latent_model_input.shape[0])
+                noise_pred = self.ace_step_transformer.decode(
+                    hidden_states=latent_model_input,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_mask=encoder_hidden_mask,
+                    output_length=latent_model_input.shape[-1],
+                    timestep=timestep,
+                ).sample
+
+            if is_repaint and i >= n_min:
+                t_i = t / 1000
+                if i + 1 < len(timesteps):
+                    t_im1 = (timesteps[i + 1]) / 1000
+                else:
+                    t_im1 = torch.zeros_like(t_i).to(self.device)
+                target_latents = target_latents.to(torch.float32)
+                prev_sample = target_latents + (t_im1 - t_i) * noise_pred
+                prev_sample = prev_sample.to(self.dtype)
+                target_latents = prev_sample
+                zt_src = (1 - t_im1) * x0 + (t_im1) * z0
+                target_latents = torch.where(repaint_mask == 1.0, target_latents, zt_src)
+            else:
+                target_latents = scheduler.step(
+                    model_output=noise_pred,
+                    timestep=t,
+                    sample=target_latents,
+                    return_dict=False,
+                    omega=omega_scale,
+                    generator=random_generators[0],
+                )[0]
+
+        if is_extend:
+            if to_right_pad_gt_latents is not None:
+                target_latents = torch.cat([target_latents, to_right_pad_gt_latents], dim=-1)
+            if to_left_pad_gt_latents is not None:
+                target_latents = torch.cat([to_right_pad_gt_latents, target_latents], dim=0)
+        return target_latents
+
+    def load_lora(self, model_with_lora_path, device="CPU"):
+        if model_with_lora_path == "none":
+            if self.ace_step_transformer_origin:
+                self.ace_step_transformer = self.ace_step_transformer_origin
+        else:
+            self.ace_step_transformer_origin = self.ace_step_transformer
+            self.update_transformer_model(model_with_lora_path, device)
+
+    def update_transformer_model(self, new_transformer_path, device="CPU"):
+        self.ace_step_transformer = OvWrapperACEStepTransformer2DModel.from_pretrained(self.core, new_transformer_path, device)

ov_dcae_decoder_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6dbfea557f63367db7d4ff51457dd2e77caaf05c5cab65597a9b3ddfc38e733
+size 158288642

ov_dcae_encoder_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eb9d38bbd0c63c7d91ab7559182d83baf6d177185e8ec03cdac757d819683cb
+size 155336722

ov_text_encoder_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c86fe17bca6108430736d10fe230a67a22b4a81bd37b72bd087f9b421f84943
+size 563722938

ov_transformer_decoder_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee39b3116fac7f25623b086e10e53f4b7679b5e93fed235d10a8561a59045200
+size 2032513000

ov_transformer_encoder_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af5049084e0350d0e1787777f6d7608735dc62ec6ef73b7cf3fd13e2c46c2a33
+size 103605115

ov_vocoder_decode_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a41d5b187bfd7ba407026f16b44a1591f33f2263803014ea1bd3b382b1d67e4
+size 205955458

ov_vocoder_mel_transform_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1f5fe2a83b10fd82b22ce3dc6dd6d46d2fdf30de7233058b3abab4fdacc060b
+size 266572

ov_vocoder_mel_transform_model.xml

@@ -0,0 +1,512 @@
+<?xml version="1.0"?>
+<net name="Model9" version="11">
+	<layers>
+		<layer id="0" name="x" type="Parameter" version="opset1">
+			<data shape="?,?" element_type="f32" />
+			<output>
+				<port id="0" precision="FP32" names="x">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="1" name="21" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="0" size="8" />
+			<output>
+				<port id="0" precision="I64" names="21" />
+			</output>
+		</layer>
+		<layer id="2" name="__module.spectrogram/aten::unsqueeze/Unsqueeze" type="Unsqueeze" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I64" />
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="23,input.1">
+					<dim>-1</dim>
+					<dim>1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="3" name="__module.spectrogram/aten::pad/Concat" type="Const" version="opset1">
+			<data element_type="i64" shape="3" offset="8" size="24" />
+			<rt_info>
+				<attribute name="precise" version="0" />
+			</rt_info>
+			<output>
+				<port id="0" precision="I64">
+					<dim>3</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="4" name="__module.spectrogram/aten::pad/ConvertLike_1_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="" offset="32" size="2" />
+			<output>
+				<port id="0" precision="FP16" />
+			</output>
+		</layer>
+		<layer id="5" name="__module.spectrogram/aten::pad/ConvertLike_1" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16" />
+			</input>
+			<output>
+				<port id="1" precision="FP32" />
+			</output>
+		</layer>
+		<layer id="6" name="__module.spectrogram/aten::pad/Pad" type="Pad" version="opset12">
+			<data pad_mode="reflect" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I64">
+					<dim>3</dim>
+				</port>
+				<port id="2" precision="I64">
+					<dim>3</dim>
+				</port>
+				<port id="3" precision="FP32" />
+			</input>
+			<output>
+				<port id="4" precision="FP32" names="25">
+					<dim>-1</dim>
+					<dim>1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="7" name="__module.spectrogram/aten::squeeze/Squeeze" type="Squeeze" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="I64" />
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="26,27,input,y">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="8" name="self.spectrogram.window_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="2048" offset="34" size="4096" />
+			<output>
+				<port id="0" precision="FP16" names="self.spectrogram.window">
+					<dim>2048</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="9" name="self.spectrogram.window" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>2048</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>2048</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="10" name="15" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="4130" size="8" />
+			<output>
+				<port id="0" precision="I64" names="15" />
+			</output>
+		</layer>
+		<layer id="11" name="14" type="Const" version="opset1">
+			<data element_type="i64" shape="" offset="4138" size="8" />
+			<output>
+				<port id="0" precision="I64" names="14" />
+			</output>
+		</layer>
+		<layer id="12" name="__module.spectrogram/aten::stft/STFT" type="STFT" version="opset15">
+			<data transpose_frames="true" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="FP32">
+					<dim>2048</dim>
+				</port>
+				<port id="2" precision="I64" />
+				<port id="3" precision="I64" />
+			</input>
+			<output>
+				<port id="4" precision="FP32" names="29,spec.3">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+					<dim>2</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="13" name="Constant_121258_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="1, 1, 1, 1" offset="4146" size="2" />
+			<output>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="14" name="Constant_121258" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="15" name="__module.spectrogram/aten::pow/Power" type="Power" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+					<dim>2</dim>
+				</port>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="30">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+					<dim>2</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="16" name="Constant_118930" type="Const" version="opset1">
+			<data element_type="i64" shape="1" offset="4148" size="8" />
+			<output>
+				<port id="0" precision="I64" names="31">
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="17" name="__module.spectrogram/aten::sum/ReduceSum" type="ReduceSum" version="opset1">
+			<data keep_dims="false" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+					<dim>2</dim>
+				</port>
+				<port id="1" precision="I64">
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="18" name="Constant_121259_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="1, 1, 1" offset="4156" size="2" />
+			<output>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="19" name="Constant_121259" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="20" name="__module.spectrogram/aten::add/Add" type="Add" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="33">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="21" name="__module.spectrogram/aten::sqrt/Sqrt" type="Sqrt" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32" names="34,35,spec,specgram">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="22" name="Transpose_121127_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="128, 1025" offset="4158" size="262400" />
+			<output>
+				<port id="0" precision="FP16">
+					<dim>128</dim>
+					<dim>1025</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="23" name="Transpose_121127" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>128</dim>
+					<dim>1025</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>128</dim>
+					<dim>1025</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="24" name="__module.mel_scale/aten::matmul/MatMul" type="MatMul" version="opset1">
+			<data transpose_a="true" transpose_b="true" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>1025</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="FP32">
+					<dim>128</dim>
+					<dim>1025</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="40">
+					<dim>-1</dim>
+					<dim>-1</dim>
+					<dim>128</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="25" name="__module.mel_scale/aten::transpose/Constant_1" type="Const" version="opset1">
+			<data element_type="i32" shape="3" offset="266558" size="12" />
+			<output>
+				<port id="0" precision="I32">
+					<dim>3</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="26" name="__module.mel_scale/aten::transpose/Transpose_1" type="Transpose" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>-1</dim>
+					<dim>128</dim>
+				</port>
+				<port id="1" precision="I32">
+					<dim>3</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="41,x_1">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="27" name="Constant_121260_compressed" type="Const" version="opset1">
+			<data element_type="f16" shape="1, 1, 1" offset="266570" size="2" />
+			<output>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="28" name="Constant_121260" type="Convert" version="opset1">
+			<data destination_type="f32" />
+			<rt_info>
+				<attribute name="decompression" version="0" />
+			</rt_info>
+			<input>
+				<port id="0" precision="FP16">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="29" name="aten::clamp/Maximum" type="Maximum" version="opset1">
+			<data auto_broadcast="numpy" />
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+				<port id="1" precision="FP32">
+					<dim>1</dim>
+					<dim>1</dim>
+					<dim>1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="2" precision="FP32" names="8">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="30" name="aten::log/Log" type="Log" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+			<output>
+				<port id="1" precision="FP32">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="31" name="Result_119110" type="Result" version="opset1">
+			<input>
+				<port id="0" precision="FP32">
+					<dim>-1</dim>
+					<dim>128</dim>
+					<dim>-1</dim>
+				</port>
+			</input>
+		</layer>
+	</layers>
+	<edges>
+		<edge from-layer="0" from-port="0" to-layer="2" to-port="0" />
+		<edge from-layer="1" from-port="0" to-layer="7" to-port="1" />
+		<edge from-layer="1" from-port="0" to-layer="2" to-port="1" />
+		<edge from-layer="2" from-port="2" to-layer="6" to-port="0" />
+		<edge from-layer="3" from-port="0" to-layer="6" to-port="1" />
+		<edge from-layer="3" from-port="0" to-layer="6" to-port="2" />
+		<edge from-layer="4" from-port="0" to-layer="5" to-port="0" />
+		<edge from-layer="5" from-port="1" to-layer="6" to-port="3" />
+		<edge from-layer="6" from-port="4" to-layer="7" to-port="0" />
+		<edge from-layer="7" from-port="2" to-layer="12" to-port="0" />
+		<edge from-layer="8" from-port="0" to-layer="9" to-port="0" />
+		<edge from-layer="9" from-port="1" to-layer="12" to-port="1" />
+		<edge from-layer="10" from-port="0" to-layer="12" to-port="2" />
+		<edge from-layer="11" from-port="0" to-layer="12" to-port="3" />
+		<edge from-layer="12" from-port="4" to-layer="15" to-port="0" />
+		<edge from-layer="13" from-port="0" to-layer="14" to-port="0" />
+		<edge from-layer="14" from-port="1" to-layer="15" to-port="1" />
+		<edge from-layer="15" from-port="2" to-layer="17" to-port="0" />
+		<edge from-layer="16" from-port="0" to-layer="17" to-port="1" />
+		<edge from-layer="17" from-port="2" to-layer="20" to-port="0" />
+		<edge from-layer="18" from-port="0" to-layer="19" to-port="0" />
+		<edge from-layer="19" from-port="1" to-layer="20" to-port="1" />
+		<edge from-layer="20" from-port="2" to-layer="21" to-port="0" />
+		<edge from-layer="21" from-port="1" to-layer="24" to-port="0" />
+		<edge from-layer="22" from-port="0" to-layer="23" to-port="0" />
+		<edge from-layer="23" from-port="1" to-layer="24" to-port="1" />
+		<edge from-layer="24" from-port="2" to-layer="26" to-port="0" />
+		<edge from-layer="25" from-port="0" to-layer="26" to-port="1" />
+		<edge from-layer="26" from-port="2" to-layer="29" to-port="0" />
+		<edge from-layer="27" from-port="0" to-layer="28" to-port="0" />
+		<edge from-layer="28" from-port="1" to-layer="29" to-port="1" />
+		<edge from-layer="29" from-port="2" to-layer="30" to-port="0" />
+		<edge from-layer="30" from-port="1" to-layer="31" to-port="0" />
+	</edges>
+	<rt_info>
+		<Runtime_version value="2025.4.0-20398-7a975177ff4-releases/2025/4" />
+		<conversion_parameters>
+			<framework value="pytorch" />
+			<is_python_object value="True" />
+		</conversion_parameters>
+	</rt_info>
+</net>