Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| import argparse | |
| import math | |
| import os | |
| from copy import deepcopy | |
| import torch | |
| from audio_diffusion.models import DiffusionAttnUnet1D | |
| from diffusion import sampling | |
| from torch import nn | |
| from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNet1DModel | |
| MODELS_MAP = { | |
| "gwf-440k": { | |
| "url": "https://model-server.zqevans2.workers.dev/gwf-440k.ckpt", | |
| "sample_rate": 48000, | |
| "sample_size": 65536, | |
| }, | |
| "jmann-small-190k": { | |
| "url": "https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt", | |
| "sample_rate": 48000, | |
| "sample_size": 65536, | |
| }, | |
| "jmann-large-580k": { | |
| "url": "https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt", | |
| "sample_rate": 48000, | |
| "sample_size": 131072, | |
| }, | |
| "maestro-uncond-150k": { | |
| "url": "https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt", | |
| "sample_rate": 16000, | |
| "sample_size": 65536, | |
| }, | |
| "unlocked-uncond-250k": { | |
| "url": "https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt", | |
| "sample_rate": 16000, | |
| "sample_size": 65536, | |
| }, | |
| "honk-140k": { | |
| "url": "https://model-server.zqevans2.workers.dev/honk-140k.ckpt", | |
| "sample_rate": 16000, | |
| "sample_size": 65536, | |
| }, | |
| } | |
| def alpha_sigma_to_t(alpha, sigma): | |
| """Returns a timestep, given the scaling factors for the clean image and for | |
| the noise.""" | |
| return torch.atan2(sigma, alpha) / math.pi * 2 | |
| def get_crash_schedule(t): | |
| sigma = torch.sin(t * math.pi / 2) ** 2 | |
| alpha = (1 - sigma**2) ** 0.5 | |
| return alpha_sigma_to_t(alpha, sigma) | |
| class Object(object): | |
| pass | |
| class DiffusionUncond(nn.Module): | |
| def __init__(self, global_args): | |
| super().__init__() | |
| self.diffusion = DiffusionAttnUnet1D(global_args, n_attn_layers=4) | |
| self.diffusion_ema = deepcopy(self.diffusion) | |
| self.rng = torch.quasirandom.SobolEngine(1, scramble=True) | |
| def download(model_name): | |
| url = MODELS_MAP[model_name]["url"] | |
| os.system(f"wget {url} ./") | |
| return f"./{model_name}.ckpt" | |
| DOWN_NUM_TO_LAYER = { | |
| "1": "resnets.0", | |
| "2": "attentions.0", | |
| "3": "resnets.1", | |
| "4": "attentions.1", | |
| "5": "resnets.2", | |
| "6": "attentions.2", | |
| } | |
| UP_NUM_TO_LAYER = { | |
| "8": "resnets.0", | |
| "9": "attentions.0", | |
| "10": "resnets.1", | |
| "11": "attentions.1", | |
| "12": "resnets.2", | |
| "13": "attentions.2", | |
| } | |
| MID_NUM_TO_LAYER = { | |
| "1": "resnets.0", | |
| "2": "attentions.0", | |
| "3": "resnets.1", | |
| "4": "attentions.1", | |
| "5": "resnets.2", | |
| "6": "attentions.2", | |
| "8": "resnets.3", | |
| "9": "attentions.3", | |
| "10": "resnets.4", | |
| "11": "attentions.4", | |
| "12": "resnets.5", | |
| "13": "attentions.5", | |
| } | |
| DEPTH_0_TO_LAYER = { | |
| "0": "resnets.0", | |
| "1": "resnets.1", | |
| "2": "resnets.2", | |
| "4": "resnets.0", | |
| "5": "resnets.1", | |
| "6": "resnets.2", | |
| } | |
| RES_CONV_MAP = { | |
| "skip": "conv_skip", | |
| "main.0": "conv_1", | |
| "main.1": "group_norm_1", | |
| "main.3": "conv_2", | |
| "main.4": "group_norm_2", | |
| } | |
| ATTN_MAP = { | |
| "norm": "group_norm", | |
| "qkv_proj": ["query", "key", "value"], | |
| "out_proj": ["proj_attn"], | |
| } | |
| def convert_resconv_naming(name): | |
| if name.startswith("skip"): | |
| return name.replace("skip", RES_CONV_MAP["skip"]) | |
| # name has to be of format main.{digit} | |
| if not name.startswith("main."): | |
| raise ValueError(f"ResConvBlock error with {name}") | |
| return name.replace(name[:6], RES_CONV_MAP[name[:6]]) | |
| def convert_attn_naming(name): | |
| for key, value in ATTN_MAP.items(): | |
| if name.startswith(key) and not isinstance(value, list): | |
| return name.replace(key, value) | |
| elif name.startswith(key): | |
| return [name.replace(key, v) for v in value] | |
| raise ValueError(f"Attn error with {name}") | |
| def rename(input_string, max_depth=13): | |
| string = input_string | |
| if string.split(".")[0] == "timestep_embed": | |
| return string.replace("timestep_embed", "time_proj") | |
| depth = 0 | |
| if string.startswith("net.3."): | |
| depth += 1 | |
| string = string[6:] | |
| elif string.startswith("net."): | |
| string = string[4:] | |
| while string.startswith("main.7."): | |
| depth += 1 | |
| string = string[7:] | |
| if string.startswith("main."): | |
| string = string[5:] | |
| # mid block | |
| if string[:2].isdigit(): | |
| layer_num = string[:2] | |
| string_left = string[2:] | |
| else: | |
| layer_num = string[0] | |
| string_left = string[1:] | |
| if depth == max_depth: | |
| new_layer = MID_NUM_TO_LAYER[layer_num] | |
| prefix = "mid_block" | |
| elif depth > 0 and int(layer_num) < 7: | |
| new_layer = DOWN_NUM_TO_LAYER[layer_num] | |
| prefix = f"down_blocks.{depth}" | |
| elif depth > 0 and int(layer_num) > 7: | |
| new_layer = UP_NUM_TO_LAYER[layer_num] | |
| prefix = f"up_blocks.{max_depth - depth - 1}" | |
| elif depth == 0: | |
| new_layer = DEPTH_0_TO_LAYER[layer_num] | |
| prefix = f"up_blocks.{max_depth - 1}" if int(layer_num) > 3 else "down_blocks.0" | |
| if not string_left.startswith("."): | |
| raise ValueError(f"Naming error with {input_string} and string_left: {string_left}.") | |
| string_left = string_left[1:] | |
| if "resnets" in new_layer: | |
| string_left = convert_resconv_naming(string_left) | |
| elif "attentions" in new_layer: | |
| new_string_left = convert_attn_naming(string_left) | |
| string_left = new_string_left | |
| if not isinstance(string_left, list): | |
| new_string = prefix + "." + new_layer + "." + string_left | |
| else: | |
| new_string = [prefix + "." + new_layer + "." + s for s in string_left] | |
| return new_string | |
| def rename_orig_weights(state_dict): | |
| new_state_dict = {} | |
| for k, v in state_dict.items(): | |
| if k.endswith("kernel"): | |
| # up- and downsample layers, don't have trainable weights | |
| continue | |
| new_k = rename(k) | |
| # check if we need to transform from Conv => Linear for attention | |
| if isinstance(new_k, list): | |
| new_state_dict = transform_conv_attns(new_state_dict, new_k, v) | |
| else: | |
| new_state_dict[new_k] = v | |
| return new_state_dict | |
| def transform_conv_attns(new_state_dict, new_k, v): | |
| if len(new_k) == 1: | |
| if len(v.shape) == 3: | |
| # weight | |
| new_state_dict[new_k[0]] = v[:, :, 0] | |
| else: | |
| # bias | |
| new_state_dict[new_k[0]] = v | |
| else: | |
| # qkv matrices | |
| trippled_shape = v.shape[0] | |
| single_shape = trippled_shape // 3 | |
| for i in range(3): | |
| if len(v.shape) == 3: | |
| new_state_dict[new_k[i]] = v[i * single_shape : (i + 1) * single_shape, :, 0] | |
| else: | |
| new_state_dict[new_k[i]] = v[i * single_shape : (i + 1) * single_shape] | |
| return new_state_dict | |
| def main(args): | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model_name = args.model_path.split("/")[-1].split(".")[0] | |
| if not os.path.isfile(args.model_path): | |
| assert ( | |
| model_name == args.model_path | |
| ), f"Make sure to provide one of the official model names {MODELS_MAP.keys()}" | |
| args.model_path = download(model_name) | |
| sample_rate = MODELS_MAP[model_name]["sample_rate"] | |
| sample_size = MODELS_MAP[model_name]["sample_size"] | |
| config = Object() | |
| config.sample_size = sample_size | |
| config.sample_rate = sample_rate | |
| config.latent_dim = 0 | |
| diffusers_model = UNet1DModel(sample_size=sample_size, sample_rate=sample_rate) | |
| diffusers_state_dict = diffusers_model.state_dict() | |
| orig_model = DiffusionUncond(config) | |
| orig_model.load_state_dict(torch.load(args.model_path, map_location=device)["state_dict"]) | |
| orig_model = orig_model.diffusion_ema.eval() | |
| orig_model_state_dict = orig_model.state_dict() | |
| renamed_state_dict = rename_orig_weights(orig_model_state_dict) | |
| renamed_minus_diffusers = set(renamed_state_dict.keys()) - set(diffusers_state_dict.keys()) | |
| diffusers_minus_renamed = set(diffusers_state_dict.keys()) - set(renamed_state_dict.keys()) | |
| assert len(renamed_minus_diffusers) == 0, f"Problem with {renamed_minus_diffusers}" | |
| assert all(k.endswith("kernel") for k in list(diffusers_minus_renamed)), f"Problem with {diffusers_minus_renamed}" | |
| for key, value in renamed_state_dict.items(): | |
| assert ( | |
| diffusers_state_dict[key].squeeze().shape == value.squeeze().shape | |
| ), f"Shape for {key} doesn't match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}" | |
| if key == "time_proj.weight": | |
| value = value.squeeze() | |
| diffusers_state_dict[key] = value | |
| diffusers_model.load_state_dict(diffusers_state_dict) | |
| steps = 100 | |
| seed = 33 | |
| diffusers_scheduler = IPNDMScheduler(num_train_timesteps=steps) | |
| generator = torch.manual_seed(seed) | |
| noise = torch.randn([1, 2, config.sample_size], generator=generator).to(device) | |
| t = torch.linspace(1, 0, steps + 1, device=device)[:-1] | |
| step_list = get_crash_schedule(t) | |
| pipe = DanceDiffusionPipeline(unet=diffusers_model, scheduler=diffusers_scheduler) | |
| generator = torch.manual_seed(33) | |
| audio = pipe(num_inference_steps=steps, generator=generator).audios | |
| generated = sampling.iplms_sample(orig_model, noise, step_list, {}) | |
| generated = generated.clamp(-1, 1) | |
| diff_sum = (generated - audio).abs().sum() | |
| diff_max = (generated - audio).abs().max() | |
| if args.save: | |
| pipe.save_pretrained(args.checkpoint_path) | |
| print("Diff sum", diff_sum) | |
| print("Diff max", diff_max) | |
| assert diff_max < 1e-3, f"Diff max: {diff_max} is too much :-/" | |
| print(f"Conversion for {model_name} successful!") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.") | |
| parser.add_argument( | |
| "--save", default=True, type=bool, required=False, help="Whether to save the converted model or not." | |
| ) | |
| parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") | |
| args = parser.parse_args() | |
| main(args) | |