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]