| from typing import Optional | |
| from nitrous_ema import PostHocEMA | |
| from omegaconf import DictConfig | |
| from mmaudio.model.networks import get_my_mmaudio | |
| def synthesize_ema(cfg: DictConfig, sigma: float, step: Optional[int]): | |
| vae = get_my_mmaudio(cfg.model) | |
| emas = PostHocEMA(vae, | |
| sigma_rels=cfg.ema.sigma_rels, | |
| update_every=cfg.ema.update_every, | |
| checkpoint_every_num_steps=cfg.ema.checkpoint_every, | |
| checkpoint_folder=cfg.ema.checkpoint_folder) | |
| synthesized_ema = emas.synthesize_ema_model(sigma_rel=sigma, step=step, device='cpu') | |
| state_dict = synthesized_ema.ema_model.state_dict() | |
| return state_dict | |