from pathlib import Path
import os,sys,json,time,threading,traceback
root=Path(__file__).parent
os.environ['ACESTEP_CHECKPOINTS_DIR']=str(root/'ace-checkpoints')
os.environ['HF_HOME']=str(root/'hf-cache')
os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
os.environ['PYTORCH_ENABLE_MPS_FALLBACK']='1'
out=root/'outputs/ace-remix';out.mkdir(exist_ok=True)
print('Importing ACE-Step official API',flush=True)
import psutil
samples=[];started=time.time();stop=False

def memory_guard():
 global stop
 p=psutil.Process()
 while not stop:
  rss=p.memory_info().rss;available=psutil.virtual_memory().available
  samples.append({'seconds':round(time.time()-started,2),'rssBytes':rss,'availableBytes':available})
  if rss>9.25*1024**3:
   (out/'memory-stop.json').write_text(json.dumps({'reason':'Own RSS exceeded 9.25 GiB; protecting shared 16 GiB machine','last':samples[-1]},indent=2));os._exit(74)
  time.sleep(2)
threading.Thread(target=memory_guard,daemon=True).start()
import torch
from acestep.handler import AceStepHandler
from acestep.inference import GenerationParams,GenerationConfig,generate_music
torch.set_num_threads(2)
class LowMemoryCoverHandler(AceStepHandler):
 def _ensure_models_present(self,*,checkpoint_path,config_path,prefer_source,vae_variant=None):
  # Cover directly conditions on source audio and never invokes the 5 Hz LM.
  # Validate every actual required component instead of triggering the default LM download.
  required=[f'{config_path}/model.safetensors',f'{config_path}/config.json',f'{config_path}/silence_latent.pt','Qwen3-Embedding-0.6B/model.safetensors','Qwen3-Embedding-0.6B/config.json','Qwen3-Embedding-0.6B/tokenizer.json','vae/diffusion_pytorch_model.safetensors','vae/config.json']
  missing=[x for x in required if not (checkpoint_path/x).is_file()]
  if missing:return 'Required core files missing: '+', '.join(missing),False
  return None
 def _load_main_model_from_checkpoint(self,**kwargs):
  # Explicit local adaptation: preserve official DiT BF16 instead of FP32 expansion on CPU.
  self.dtype=torch.bfloat16
  return super()._load_main_model_from_checkpoint(**kwargs)
params=GenerationParams(task_type='cover',src_audio=str(root/'inputs/original-eight-second-melody.wav'),caption='An original warm acoustic instrumental, plucked acoustic guitar carrying the clear C major motif, gentle bass and very light brushed percussion, no vocals, intimate dry studio sound.',lyrics='[Instrumental]',instrumental=True,bpm=60,keyscale='C Major',timesignature='4',duration=8.0,inference_steps=4,guidance_scale=1.0,seed=20261024,audio_cover_strength=0.6,thinking=False,use_cot_metas=False,use_cot_caption=False,use_cot_language=False,dcw_enabled=False)
config=GenerationConfig(batch_size=1,use_random_seed=False,seeds=[20261024],audio_format='wav')
receipt={'tool':'ACE-Step 1.5 official generate_music API','sourceCommit':'ca1e85fe9430179831e6bc6be790c332190a3866','modelRevision':'19671f406d603126926c1b7e2adc169acbcade22','runtimeAdaptations':['CPU DiT BF16 override to avoid FP32 memory expansion; model/algorithm unchanged','Validate only actual three core components for source-conditioned Cover; LM not downloaded or initialized','CPU two threads; offload flags enabled; no MLX/compile/quantization; own RSS stop at 9.25 GiB'],'params':params.to_dict() if hasattr(params,'to_dict') else vars(params),'config':config.to_dict(),'sourceIdentity':'Program-synthesized original eight-second motif; not human humming','success':False}
try:
 handler=LowMemoryCoverHandler()
 print('Initializing three core components; LM disabled',flush=True)
 status,success=handler.initialize_service(project_root=str(root/'ace-source'),config_path='acestep-v15-turbo',device='cpu',offload_to_cpu=True,offload_dit_to_cpu=True,use_mlx_dit=False,compile_model=False,quantization=None)
 receipt['initialization']={'success':success,'status':status,'elapsedSeconds':time.time()-started}
 print('INITIALIZATION',success,status,flush=True)
 if not success:raise RuntimeError(status)
 result=generate_music(handler,None,params=params,config=config,save_dir=str(out))
 receipt['success']=bool(result.success)
 receipt['statusMessage']=result.status_message
 receipt['audios']=[{k:v for k,v in x.items() if k not in ['tensor','audio','waveform']} for x in result.audios]
 print('RESULT',receipt['success'],receipt['statusMessage'],flush=True)
except Exception as exc:
 receipt['error']={'type':type(exc).__name__,'message':str(exc)};traceback.print_exc()
finally:
 stop=True;receipt['elapsedSeconds']=time.time()-started;receipt['maxRssBytes']=max((x['rssBytes'] for x in samples),default=0)
 (out/'memory-samples.json').write_text(json.dumps(samples,indent=2))
 (out/'execution.json').write_text(json.dumps(receipt,indent=2,default=str))
 print('EXECUTION_RECEIPT_WRITTEN',flush=True)