Upload conversion_script.py

conversion_script.py

@@ -0,0 +1,233 @@
+import json
+import os
+from contextlib import nullcontext
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoTokenizer, T5EncoderModel
+
+from diffusers import (
+    AutoencoderOobleck,
+    CosineDPMSolverMultistepScheduler,
+    StableAudioDiTModel,
+    StableAudioPipeline,
+    StableAudioProjectionModel,
+)
+from diffusers.models.model_loading_utils import load_model_dict_into_meta
+from diffusers.utils import is_accelerate_available
+
+if is_accelerate_available():
+    from accelerate import init_empty_weights
+
+# ==========================================
+# HARDCODED PATHS FOR ENVIRONMENT
+# ==========================================
+CHECKPOINT_PATH = r"\Foundation_1.safetensors" # YOU MUST DOWNLOAD THIS
+CONFIG_PATH = r"\model_config.json" # YOU MUST DOWNLOAD THIS
+SAVE_DIRECTORY = r"\foundation_diffusers"
+
+device = "cpu"
+dtype = torch.float32
+
+# ==========================================
+
+def convert_stable_audio_state_dict_to_diffusers(state_dict, num_autoencoder_layers=5):
+    projection_model_state_dict = {
+        k.replace("conditioner.conditioners.", "").replace("embedder.embedding", "time_positional_embedding"): v
+        for (k, v) in state_dict.items()
+        if "conditioner.conditioners" in k
+    }
+
+    for key, value in list(projection_model_state_dict.items()):
+        new_key = key.replace("seconds_start", "start_number_conditioner").replace(
+            "seconds_total", "end_number_conditioner"
+        )
+        projection_model_state_dict[new_key] = projection_model_state_dict.pop(key)
+
+    model_state_dict = {k.replace("model.model.", ""): v for (k, v) in state_dict.items() if "model.model." in k}
+    for key, value in list(model_state_dict.items()):
+        new_key = (
+            key.replace("transformer.", "")
+            .replace("layers", "transformer_blocks")
+            .replace("self_attn", "attn1")
+            .replace("cross_attn", "attn2")
+            .replace("ff.ff", "ff.net")
+        )
+        new_key = (
+            new_key.replace("pre_norm", "norm1")
+            .replace("cross_attend_norm", "norm2")
+            .replace("ff_norm", "norm3")
+            .replace("to_out", "to_out.0")
+        )
+        new_key = new_key.replace("gamma", "weight").replace("beta", "bias") 
+
+        new_key = (
+            new_key.replace("project", "proj")
+            .replace("to_timestep_embed", "timestep_proj")
+            .replace("timestep_features", "time_proj")
+            .replace("to_global_embed", "global_proj")
+            .replace("to_cond_embed", "cross_attention_proj")
+        )
+
+        if new_key == "time_proj.weight":
+            model_state_dict[key] = model_state_dict[key].squeeze(1)
+
+        if "to_qkv" in new_key:
+            q, k, v = torch.chunk(model_state_dict.pop(key), 3, dim=0)
+            model_state_dict[new_key.replace("qkv", "q")] = q
+            model_state_dict[new_key.replace("qkv", "k")] = k
+            model_state_dict[new_key.replace("qkv", "v")] = v
+        elif "to_kv" in new_key:
+            k, v = torch.chunk(model_state_dict.pop(key), 2, dim=0)
+            model_state_dict[new_key.replace("kv", "k")] = k
+            model_state_dict[new_key.replace("kv", "v")] = v
+        else:
+            model_state_dict[new_key] = model_state_dict.pop(key)
+
+    autoencoder_state_dict = {
+        k.replace("pretransform.model.", "").replace("coder.layers.0", "coder.conv1"): v
+        for (k, v) in state_dict.items()
+        if "pretransform.model." in k
+    }
+
+    for key, _ in list(autoencoder_state_dict.items()):
+        new_key = key
+        if "coder.layers" in new_key:
+            idx = int(new_key.split("coder.layers.")[1].split(".")[0])
+            new_key = new_key.replace(f"coder.layers.{idx}", f"coder.block.{idx - 1}")
+
+            if "encoder" in new_key:
+                for i in range(3):
+                    new_key = new_key.replace(f"block.{idx - 1}.layers.{i}", f"block.{idx - 1}.res_unit{i + 1}")
+                new_key = new_key.replace(f"block.{idx - 1}.layers.3", f"block.{idx - 1}.snake1")
+                new_key = new_key.replace(f"block.{idx - 1}.layers.4", f"block.{idx - 1}.conv1")
+            else:
+                for i in range(2, 5):
+                    new_key = new_key.replace(f"block.{idx - 1}.layers.{i}", f"block.{idx - 1}.res_unit{i - 1}")
+                new_key = new_key.replace(f"block.{idx - 1}.layers.0", f"block.{idx - 1}.snake1")
+                new_key = new_key.replace(f"block.{idx - 1}.layers.1", f"block.{idx - 1}.conv_t1")
+
+            new_key = new_key.replace("layers.0.beta", "snake1.beta")
+            new_key = new_key.replace("layers.0.alpha", "snake1.alpha")
+            new_key = new_key.replace("layers.2.beta", "snake2.beta")
+            new_key = new_key.replace("layers.2.alpha", "snake2.alpha")
+            new_key = new_key.replace("layers.1.bias", "conv1.bias")
+            new_key = new_key.replace("layers.1.weight_", "conv1.weight_")
+            new_key = new_key.replace("layers.3.bias", "conv2.bias")
+            new_key = new_key.replace("layers.3.weight_", "conv2.weight_")
+
+            if idx == num_autoencoder_layers + 1:
+                new_key = new_key.replace(f"block.{idx - 1}", "snake1")
+            elif idx == num_autoencoder_layers + 2:
+                new_key = new_key.replace(f"block.{idx - 1}", "conv2")
+
+        value = autoencoder_state_dict.pop(key)
+        if "snake" in new_key:
+            value = value.unsqueeze(0).unsqueeze(-1)
+        if new_key in autoencoder_state_dict:
+            raise ValueError(f"{new_key} already in state dict.")
+        autoencoder_state_dict[new_key] = value
+
+    return model_state_dict, projection_model_state_dict, autoencoder_state_dict
+
+print("Reading config...")
+with open(CONFIG_PATH) as f_in:
+    config_dict = json.load(f_in)
+
+conditioning_dict = {
+    conditioning["id"]: conditioning["config"] for conditioning in config_dict["model"]["conditioning"]["configs"]
+}
+
+t5_model_config = conditioning_dict["prompt"]
+
+print("Downloading/Loading T5 text encoder...")
+text_encoder = T5EncoderModel.from_pretrained(t5_model_config["t5_model_name"])
+tokenizer = AutoTokenizer.from_pretrained(
+    t5_model_config["t5_model_name"], truncation=True, model_max_length=t5_model_config["max_length"]
+)
+
+scheduler = CosineDPMSolverMultistepScheduler(
+    sigma_min=0.3,
+    sigma_max=500,
+    solver_order=2,
+    prediction_type="v_prediction",
+    sigma_data=1.0,
+    sigma_schedule="exponential",
+)
+ctx = init_empty_weights if is_accelerate_available() else nullcontext
+
+print("Loading SafeTensors checkpoint...")
+orig_state_dict = load_file(CHECKPOINT_PATH, device=device)
+
+model_config = config_dict["model"]["diffusion"]["config"]
+
+print("Converting weights (this might take a moment)...")
+model_state_dict, projection_model_state_dict, autoencoder_state_dict = convert_stable_audio_state_dict_to_diffusers(
+    orig_state_dict
+)
+
+print("Building Models...")
+with ctx():
+    projection_model = StableAudioProjectionModel(
+        text_encoder_dim=text_encoder.config.d_model,
+        conditioning_dim=config_dict["model"]["conditioning"]["cond_dim"],
+        min_value=conditioning_dict["seconds_start"]["min_val"],
+        max_value=conditioning_dict["seconds_start"]["max_val"],
+    )
+if is_accelerate_available():
+    load_model_dict_into_meta(projection_model, projection_model_state_dict)
+else:
+    projection_model.load_state_dict(projection_model_state_dict)
+
+attention_head_dim = model_config["embed_dim"] // model_config["num_heads"]
+with ctx():
+    model = StableAudioDiTModel(
+        sample_size=int(config_dict["sample_size"])
+        / int(config_dict["model"]["pretransform"]["config"]["downsampling_ratio"]),
+        in_channels=model_config["io_channels"],
+        num_layers=model_config["depth"],
+        attention_head_dim=attention_head_dim,
+        num_key_value_attention_heads=model_config["cond_token_dim"] // attention_head_dim,
+        num_attention_heads=model_config["num_heads"],
+        out_channels=model_config["io_channels"],
+        cross_attention_dim=model_config["cond_token_dim"],
+        time_proj_dim=256,
+        global_states_input_dim=model_config["global_cond_dim"],
+        cross_attention_input_dim=model_config["cond_token_dim"],
+    )
+if is_accelerate_available():
+    load_model_dict_into_meta(model, model_state_dict)
+else:
+    model.load_state_dict(model_state_dict)
+
+autoencoder_config = config_dict["model"]["pretransform"]["config"]
+with ctx():
+    autoencoder = AutoencoderOobleck(
+        encoder_hidden_size=autoencoder_config["encoder"]["config"]["channels"],
+        downsampling_ratios=autoencoder_config["encoder"]["config"]["strides"],
+        decoder_channels=autoencoder_config["decoder"]["config"]["channels"],
+        decoder_input_channels=autoencoder_config["decoder"]["config"]["latent_dim"],
+        audio_channels=autoencoder_config["io_channels"],
+        channel_multiples=autoencoder_config["encoder"]["config"]["c_mults"],
+        sampling_rate=config_dict["sample_rate"],
+    )
+
+if is_accelerate_available():
+    load_model_dict_into_meta(autoencoder, autoencoder_state_dict)
+else:
+    autoencoder.load_state_dict(autoencoder_state_dict)
+
+print("Saving final diffusers pipeline...")
+os.makedirs(SAVE_DIRECTORY, exist_ok=True)
+pipeline = StableAudioPipeline(
+    transformer=model,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+    vae=autoencoder,
+    projection_model=projection_model,
+)
+
+pipeline.to(dtype).save_pretrained(SAVE_DIRECTORY)
+
+print(f"✅ DONE! Pipeline successfully saved to {SAVE_DIRECTORY}")