diffsynth 2.1.6__tar.gz → 2.1.7__tar.gz

{diffsynth-2.1.6 → diffsynth-2.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.1.6
+Version: 2.1.7
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License-Expression: Apache-2.0
@@ -43,6 +43,8 @@ Provides-Extra: npu-aarch64
 Requires-Dist: torch==2.7.1; extra == "npu-aarch64"
 Requires-Dist: torch-npu==2.7.1; extra == "npu-aarch64"
 Requires-Dist: torchvision==0.22.1; extra == "npu-aarch64"
+Provides-Extra: controlnet
+Requires-Dist: controlnet_aux; extra == "controlnet"
 Provides-Extra: infiniteyou
 Requires-Dist: insightface; extra == "infiniteyou"
 Requires-Dist: facexlib; extra == "infiniteyou"
@@ -397,6 +399,7 @@ Model overview:
     - LTX-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html), [Example code](/examples/ltx2/)
     - Wan: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html), [Example code](/examples/wanvideo/)
 - Audio generation
+    - DiffSynth-Music: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/DiffSynth-Music.html), [Example code](/examples/diffsynth_music/)
     - MiniMax-Music3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-Music3.html), [Example code](/examples/minimax_music3/)
     - ACE-Step: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/ACE-Step.html), [Example code](/examples/ace_step/)
 - Evaluation models: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Image-Quality-Metrics.html), [Example code](/examples/image_quality_metric/)
@@ -651,6 +654,7 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44
 | SenseNova-U1 | [SenseNova/SenseNova-U1.5-8B-MoT-SFT: T2I](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT-SFT) | [code](/examples/sensenova_u1/model_inference/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_inference_low_vram/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_training/full/SenseNova-U1.5-8B-MoT-SFT.sh) | [code](/examples/sensenova_u1/model_training/validate_full/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_training/lora/SenseNova-U1.5-8B-MoT-SFT.sh) | [code](/examples/sensenova_u1/model_training/validate_lora/SenseNova-U1.5-8B-MoT-SFT.py) |
 | SenseNova-U1 | [SenseNova/SenseNova-U1.5-8B-MoT-SFT: Edit](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT-SFT) | [code](/examples/sensenova_u1/model_inference/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_inference_low_vram/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_training/full/SenseNova-U1.5-8B-MoT-SFT-Edit.sh) | [code](/examples/sensenova_u1/model_training/validate_full/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_training/lora/SenseNova-U1.5-8B-MoT-SFT-Edit.sh) | [code](/examples/sensenova_u1/model_training/validate_lora/SenseNova-U1.5-8B-MoT-SFT-Edit.py) |
 | MiniMax-Music3 | [MiniMax/MiniMax-Music3](https://www.modelscope.cn/models/MiniMax/MiniMax-Music3) | [code](/examples/minimax_music3/model_inference/MiniMax-Music3.py) | [code](/examples/minimax_music3/model_inference_low_vram/MiniMax-Music3.py) | — | — | — | — |
+| DiffSynth-Music | [DiffSynth-Studio/DiffSynth-Music](https://www.modelscope.cn/models/DiffSynth-Studio/DiffSynth-Music) | [code](/examples/diffsynth_music/model_inference/DiffSynth-Music.py) | [code](/examples/diffsynth_music/model_inference_low_vram/DiffSynth-Music.py) | [code](/examples/diffsynth_music/model_training/full/DiffSynth-Music.sh) | [code](/examples/diffsynth_music/model_training/validate_full/DiffSynth-Music.py) | - | - |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: Ref2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Ref2VA.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: Retake](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Retake.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Retake.py) | - | - | - | - |

{diffsynth-2.1.6 → diffsynth-2.1.7}/README.md

@@ -330,6 +330,7 @@ Model overview:
     - LTX-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html), [Example code](/examples/ltx2/)
     - Wan: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html), [Example code](/examples/wanvideo/)
 - Audio generation
+    - DiffSynth-Music: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/DiffSynth-Music.html), [Example code](/examples/diffsynth_music/)
     - MiniMax-Music3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-Music3.html), [Example code](/examples/minimax_music3/)
     - ACE-Step: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/ACE-Step.html), [Example code](/examples/ace_step/)
 - Evaluation models: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Image-Quality-Metrics.html), [Example code](/examples/image_quality_metric/)
@@ -584,6 +585,7 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44
 | SenseNova-U1 | [SenseNova/SenseNova-U1.5-8B-MoT-SFT: T2I](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT-SFT) | [code](/examples/sensenova_u1/model_inference/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_inference_low_vram/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_training/full/SenseNova-U1.5-8B-MoT-SFT.sh) | [code](/examples/sensenova_u1/model_training/validate_full/SenseNova-U1.5-8B-MoT-SFT.py) | [code](/examples/sensenova_u1/model_training/lora/SenseNova-U1.5-8B-MoT-SFT.sh) | [code](/examples/sensenova_u1/model_training/validate_lora/SenseNova-U1.5-8B-MoT-SFT.py) |
 | SenseNova-U1 | [SenseNova/SenseNova-U1.5-8B-MoT-SFT: Edit](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT-SFT) | [code](/examples/sensenova_u1/model_inference/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_inference_low_vram/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_training/full/SenseNova-U1.5-8B-MoT-SFT-Edit.sh) | [code](/examples/sensenova_u1/model_training/validate_full/SenseNova-U1.5-8B-MoT-SFT-Edit.py) | [code](/examples/sensenova_u1/model_training/lora/SenseNova-U1.5-8B-MoT-SFT-Edit.sh) | [code](/examples/sensenova_u1/model_training/validate_lora/SenseNova-U1.5-8B-MoT-SFT-Edit.py) |
 | MiniMax-Music3 | [MiniMax/MiniMax-Music3](https://www.modelscope.cn/models/MiniMax/MiniMax-Music3) | [code](/examples/minimax_music3/model_inference/MiniMax-Music3.py) | [code](/examples/minimax_music3/model_inference_low_vram/MiniMax-Music3.py) | — | — | — | — |
+| DiffSynth-Music | [DiffSynth-Studio/DiffSynth-Music](https://www.modelscope.cn/models/DiffSynth-Studio/DiffSynth-Music) | [code](/examples/diffsynth_music/model_inference/DiffSynth-Music.py) | [code](/examples/diffsynth_music/model_inference_low_vram/DiffSynth-Music.py) | [code](/examples/diffsynth_music/model_training/full/DiffSynth-Music.sh) | [code](/examples/diffsynth_music/model_training/validate_full/DiffSynth-Music.py) | - | - |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: Ref2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Ref2VA.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: Retake](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Retake.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Retake.py) | - | - | - | - |

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/configs/model_configs.py

@@ -1202,11 +1202,24 @@ ace_step_series = [
         "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_tokenizer.AceStepTokenizerStateDictConverter",
     },
     {
-        # Example: ???
+        # Example: ModelConfig(model_id="DiffSynth-Studio/DiffSynth-Music-Tools", origin_file_pattern="track_separator/model.safetensors")
         "model_hash": "ff74b1806e6a0b52e7bbd1d3df2d26d1",
         "model_name": "demucs",
         "model_class": "diffsynth.models.demucs.HTDemucs",
     },
+    {
+        # Example: ModelConfig(model_id="DiffSynth-Studio/DiffSynth-Music-Tools", origin_file_pattern="transformer/model.safetensors")
+        "model_hash": "167b8d453605e3d72076f88790528f8e",
+        "model_name": "diffsynth_music_dit",
+        "model_class": "diffsynth.models.diffsynth_music_dit.DiffSynthMusicDiTModel",
+    },
+    {
+        # Example: ModelConfig(model_id="DiffSynth-Studio/DiffSynth-Music-Tools", origin_file_pattern="conditioner/model.safetensors")
+        "model_hash": "d6c90aec3f282bde16298e0221413098",
+        "model_name": "ace_step_conditioner",
+        "model_class": "diffsynth.models.ace_step_conditioner.AceStepConditionEncoder",
+        "extra_kwargs": {"placeholder_shape": (1, 15000, 64)},
+    },
 ]
 
 image_metrics_series = [

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/configs/vram_management_module_maps.py

@@ -80,7 +80,7 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
         "diffsynth.models.qwen_image_vae.QwenImageRMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
-    "diffsynth.models.qwen_image_controlnet.BlockWiseControlBlock": {
+    "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet": {
         "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
     },
@@ -339,7 +339,7 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
         "transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
-    "diffsynth.models.joyai_image_dit.Transformer3DModel": {
+    "diffsynth.models.joyai_image_dit.JoyAIImageDiT": {
         "diffsynth.models.joyai_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
         "diffsynth.models.joyai_image_dit.ModulateWan": "diffsynth.core.vram.layers.AutoWrappedModule",
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
@@ -512,6 +512,27 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
         "diffsynth.models.sensenova_u1_dit.SenseNovaU1RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
+    "diffsynth.models.diffsynth_music_dit.DiffSynthMusicDiTModel": {
+        "diffsynth.models.diffsynth_music_dit.DiffSynthMusicDiTLayer": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "transformers.models.qwen3.modeling_qwen3.Qwen3RotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.demucs.HTDemucs": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.ConvTranspose2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.TransformerEncoderLayer": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.demucs.ScaledEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.demucs.LayerScale": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.MultiheadAttention": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
 }
 
 def QwenImageTextEncoder_Module_Map_Updater():

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/core/data/operators.py

@@ -339,3 +339,48 @@ class LoadPureAudioWithTorchaudio(DataProcessingOperator):
         except Exception as e:
             print(f"Cannot load audio in {data} due to {e}. The audio will be `None`.")
             return None
+
+
+class LoadMultiTrackAudio(DataProcessingOperator):
+    def __init__(self, target_sample_rate=48000, max_audio_duration=None, division_factor=1):
+        self.target_sample_rate = target_sample_rate
+        self.max_audio_duration = max_audio_duration
+        self.division_factor = division_factor
+        import torchaudio
+        self.audio_loader = torchaudio.load
+        self.audio_resampler = torchaudio.functional.resample
+
+    def load_audio(self, path):
+        waveform, sample_rate = self.audio_loader(path)
+        if len(waveform.shape) == 2 and waveform.shape[0] == 1:
+            waveform = repeat(waveform, "c l -> (n c) l", n=2)
+        if self.target_sample_rate is not None and sample_rate != self.target_sample_rate:
+            waveform = self.audio_resampler(waveform, sample_rate, self.target_sample_rate)
+            sample_rate = self.target_sample_rate
+        if self.max_audio_duration is not None and waveform.shape[1] > sample_rate * self.max_audio_duration:
+            waveform = waveform[:, :int(sample_rate * self.max_audio_duration)]
+        return waveform
+
+    def load_latents(self, path):
+        latents = torch.load(path, weights_only=True, map_location="cpu")
+        return latents
+
+    def load_single_data(self, path):
+        if path is None:
+            return None
+        elif path.endswith(".pth"):
+            return self.load_latents(path)
+        else:
+            return self.load_audio(path)
+
+    def __call__(self, data):
+        if isinstance(data, str):
+            return self.load_single_data(data)
+        else:
+            audio = {}
+            for name, path in data.items():
+                audio[name] = self.load_single_data(path)
+            min_length = min([audio[name].shape[1] for name in audio if audio[name] is not None])
+            min_length = min_length // self.division_factor * self.division_factor
+            audio = {name: audio[name][:, :min_length] for name in audio}
+        return audio

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/diffusion/flow_match.py

@@ -208,6 +208,7 @@ class FlowMatchScheduler():
             for timestep in target_timesteps:
                 timestep_id = torch.argmin((timesteps - timestep).abs())
                 timesteps[timestep_id] = timestep
+                sigmas[timestep_id] = timestep / num_train_timesteps
         return sigmas, timesteps
 
     @staticmethod

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/diffusion/runner.py

@@ -30,6 +30,26 @@ def save_training_args(args):
         print(f"Warning: failed to save training arguments: {e}")
 
 
+def exclude_quantized_params_from_ddp_sync(accelerator: Accelerator, model: DiffusionTrainingModule):
+    """DDP broadcasts every parameter when it is constructed, but a quantized weight backed by a
+    tensor subclass cannot be flattened into a broadcast bucket. Such weights are frozen and every
+    rank loads them from the same checkpoint, so let DDP skip them."""
+    try:
+        from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+        quant_configs = [module.quantize_config for module in model.modules() if getattr(module, "quantize_config", None) is not None]
+        ignored = [
+            f"{name}.weight" for name, module in model.named_modules()
+            if any(quantize.is_quantized_linear(module) for quantize in quant_configs)
+            and not module.weight.requires_grad and is_traceable_wrapper_subclass(module.weight)
+        ]
+        if len(ignored) > 0:
+            model._ddp_params_and_buffers_to_ignore = ignored
+            if accelerator.is_main_process:
+                print(f"{len(ignored)} quantized weights are excluded from DDP state synchronization.")
+    except Exception as e:
+        print(f"Warning: failed to exclude quantized weights from DDP state synchronization: {e}")
+
+
 def launch_training_task(
     accelerator: Accelerator,
     dataset: torch.utils.data.Dataset,
@@ -72,6 +92,7 @@ def launch_training_task(
         offload_manager = OffloadTrainingManager(model, accelerator.device, enable_optimizer_cpu_offload, cpu_offload_split_threshold)
     else:
         model.to(device=accelerator.device)
+        exclude_quantized_params_from_ddp_sync(accelerator, model)
         model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
 
     initialize_deepspeed_gradient_checkpointing(accelerator)
@@ -117,6 +138,7 @@ def launch_data_process_task(
         model.pipe.device = accelerator.device
     else:
         model.to(device=accelerator.device)
+        exclude_quantized_params_from_ddp_sync(accelerator, model)
         model, dataloader = accelerator.prepare(model, dataloader)
     
     for data_id, data in enumerate(tqdm(dataloader)):

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/diffusion/template.py

@@ -31,7 +31,7 @@ def check_template_model_format(model):
         raise NotImplementedError("`**kwargs` is not included in `forward`.")
 
 
-def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose=1):
+def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose=1, state_dict=None):
     spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py"))
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
@@ -44,6 +44,7 @@ def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose
             path=os.path.join(path, getattr(module, 'TEMPLATE_MODEL_PATH')),
             torch_dtype=torch_dtype,
             device=device,
+            state_dict=state_dict,
         )
     else:
         # Without `TEMPLATE_MODEL_PATH`, a randomly initialized model or a non-model module will be loaded.
@@ -97,7 +98,7 @@ class TemplatePipeline(torch.nn.Module):
             for model_config in model_configs:
                 TemplatePipeline.check_vram_config(model_config)
                 model_config.download_if_necessary()
-                model = load_template_model(model_config.path, torch_dtype=torch_dtype, device=device)
+                model = load_template_model(model_config.path, torch_dtype=torch_dtype, device=device, state_dict=model_config.state_dict)
                 models.append(model)
             self.models = torch.nn.ModuleList(models)
 
@@ -164,7 +165,7 @@ class TemplatePipeline(torch.nn.Module):
         if self.lazy_loading:
             model_config = self.model_configs[model_id]
             model_config.download_if_necessary()
-            model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device)
+            model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device, state_dict=model_config.state_dict)
         else:
             model = self.models[model_id]
         return model

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/models/ace_step_conditioner.py

@@ -585,6 +585,7 @@ class AceStepConditionEncoder(nn.Module):
         timbre_hidden_dim: int = 64,
         num_lyric_encoder_hidden_layers: int = 8,
         num_timbre_encoder_hidden_layers: int = 4,
+        placeholder_shape=None,
         **kwargs,
     ):
         super().__init__()
@@ -648,6 +649,8 @@ class AceStepConditionEncoder(nn.Module):
             timbre_hidden_dim=timbre_hidden_dim,
             num_timbre_encoder_hidden_layers=num_timbre_encoder_hidden_layers,
         )
+        if placeholder_shape is not None:
+            self.placeholder_audio = torch.nn.Parameter(torch.zeros(placeholder_shape))
 
     def forward(
         self,
@@ -658,6 +661,9 @@ class AceStepConditionEncoder(nn.Module):
         reference_latents: Optional[torch.Tensor] = None,
         refer_audio_order_mask: Optional[torch.LongTensor] = None,
     ):
+        if reference_latents is None:
+            reference_latents = self.placeholder_audio[:, :750, :].to(device=text_hidden_states.device, dtype=text_hidden_states.dtype)
+            refer_audio_order_mask = torch.tensor([0], device=text_hidden_states.device, dtype=torch.long)
         text_hidden_states = self.text_projector(text_hidden_states)
         lyric_encoder_outputs = self.lyric_encoder(
             inputs_embeds=lyric_hidden_states,

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/models/ace_step_vae.py

@@ -196,13 +196,13 @@ class OobleckDiagonalGaussianDistribution(object):
 
     def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
         # make sure sample is on the same device as the parameters and has same dtype
-        sample = torch.randn(
-            self.mean.shape,
-            generator=generator,
-            device=self.parameters.device,
-            dtype=self.parameters.dtype,
-        )
-        x = self.mean + self.std * sample
+        # sample = torch.randn(
+        #     self.mean.shape,
+        #     generator=generator,
+        #     device=self.parameters.device,
+        #     dtype=self.parameters.dtype,
+        # )
+        x = self.mean
         return x
 
     def kl(self, other: "OobleckDiagonalGaussianDistribution" = None) -> torch.Tensor:

{diffsynth-2.1.6 → diffsynth-2.1.7}/diffsynth/models/demucs.py

@@ -400,7 +400,9 @@ class HTDemucs(nn.Module):
         out = out * (ref.std() + 1e-8) + ref.mean()
         out = out / max(1.01 * out.abs().max(), 1)
         out = out.clamp_(-1, 1).cpu()
-        out = out[self.sources.index(track)]
+        if not isinstance(track, list):
+            track = [track]
+        out = sum([out[self.sources.index(t)] for t in track], 0)
         return out
 
 def center_trim(tensor, length):

diffsynth-2.1.7/diffsynth/models/diffsynth_music_dit.py

@@ -0,0 +1,237 @@
+import math, torch
+from ..core import gradient_checkpoint_forward
+from ..core.attention.attention import attention_forward
+from einops import rearrange
+from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
+
+
+class TimestepEmbedding(torch.nn.Module):
+    def __init__(self, in_channels, time_embed_dim, scale=1):
+        super().__init__()
+        self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, bias=True)
+        self.act1 = torch.nn.SiLU()
+        self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim, bias=True)
+        self.in_channels = in_channels
+        self.act2 = torch.nn.SiLU()
+        self.time_proj = torch.nn.Linear(time_embed_dim, time_embed_dim * 6)
+        self.scale = scale
+
+    def timestep_embedding(self, t, dim, max_period=10000):
+        t = t * self.scale
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.in_channels)
+        temb = self.linear_1(t_freq.to(t.dtype))
+        temb = self.act1(temb)
+        temb = self.linear_2(temb)
+        timestep_proj = self.time_proj(self.act2(temb)).unflatten(1, (6, -1))
+        return temb, timestep_proj
+
+
+class DiffSynthMusicTimestepEmbedding(torch.nn.Module):
+    def __init__(self, in_channels, time_embed_dim):
+        super().__init__()
+        self.time_embed = TimestepEmbedding(in_channels, time_embed_dim)
+        self.time_embed_r = TimestepEmbedding(in_channels, time_embed_dim)
+
+    def forward(self, timestep):
+        timestep_r = timestep
+        temb_t, timestep_proj_t = self.time_embed(timestep)
+        temb_r, timestep_proj_r = self.time_embed_r(timestep - timestep_r)
+        temb = temb_t + temb_r
+        timestep_proj = timestep_proj_t + timestep_proj_r
+        return temb, timestep_proj
+
+
+class DiffSynthMusicAttention(torch.nn.Module):
+    def __init__(self, dim, num_heads_q, num_heads_kv, head_dim):
+        super().__init__()
+        self.head_dim = head_dim
+        self.q_proj = torch.nn.Linear(dim, num_heads_q * self.head_dim, bias=False)
+        self.k_proj = torch.nn.Linear(dim, num_heads_kv * self.head_dim, bias=False)
+        self.v_proj = torch.nn.Linear(dim, num_heads_kv * self.head_dim, bias=False)
+        self.o_proj = torch.nn.Linear(num_heads_q * self.head_dim, dim, bias=False)
+        self.q_norm = torch.nn.RMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = torch.nn.RMSNorm(self.head_dim, eps=1e-6)
+
+    def rotate_half(self, x):
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rotary_pos_emb(self, q, k, cos, sin, unsqueeze_dim=2):
+        cos = cos.unsqueeze(unsqueeze_dim)
+        sin = sin.unsqueeze(unsqueeze_dim)
+        q_embed = (q * cos) + (self.rotate_half(q) * sin)
+        k_embed = (k * cos) + (self.rotate_half(k) * sin)
+        return q_embed, k_embed
+
+    def forward(self, x, y=None, window_size=None, pos_emb=None, return_kv=False, kv_cache=None) -> torch.Tensor:
+        q = self.q_proj(x)
+        q = rearrange(q, "b s (n d) -> b s n d", d=self.head_dim)
+        q = self.q_norm(q)
+
+        if y is None: y = x
+        k, v = self.k_proj(y), self.v_proj(y)
+        k, v = rearrange(k, "b s (n d) -> b s n d", d=self.head_dim), rearrange(v, "b s (n d) -> b s n d", d=self.head_dim)
+        k = self.k_norm(k)
+
+        if pos_emb is not None: q, k = self.apply_rotary_pos_emb(q, k, *pos_emb)
+        if kv_cache is not None:
+            k = torch.concat([k, kv_cache[0]], dim=1)
+            v = torch.concat([v, kv_cache[1]], dim=1)
+        attn_output = attention_forward(
+            q, k, v,
+            q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s (n d)",
+            window_size=window_size,
+        )
+        attn_output = self.o_proj(attn_output)
+        if return_kv:
+            return attn_output, (k, v)
+        else:
+            return attn_output
+
+
+class MLP(torch.nn.Module):
+    def __init__(self, dim, dim_hidden):
+        super().__init__()
+        self.gate_proj = torch.nn.Linear(dim, dim_hidden, bias=False)
+        self.up_proj = torch.nn.Linear(dim, dim_hidden, bias=False)
+        self.down_proj = torch.nn.Linear(dim_hidden, dim, bias=False)
+        self.act_fn = torch.nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DiffSynthMusicDiTLayer(torch.nn.Module):
+    def __init__(self, dim=2560, num_heads_q=32, num_heads_kv=8, head_dim=128, dim_mlp=9728, window_size=None):
+        super().__init__()
+        self.self_attn_norm = torch.nn.RMSNorm(dim, eps=1e-6)
+        self.self_attn = DiffSynthMusicAttention(dim=dim, num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, head_dim=head_dim)
+        self.cross_attn_norm = torch.nn.RMSNorm(dim, eps=1e-6)
+        self.cross_attn = DiffSynthMusicAttention(dim=dim, num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, head_dim=head_dim)
+        self.mlp_norm = torch.nn.RMSNorm(dim, eps=1e-6)
+        self.mlp = MLP(dim, dim_mlp)
+        self.scale_shift_table = torch.nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.window_size = window_size
+
+    def forward(self, x, y, pos_emb, temb, return_kv=False, kv_cache=None) -> torch.Tensor:
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (self.scale_shift_table.to(dtype=x.dtype, device=x.device) + temb).chunk(6, dim=1)
+        x_hidden = self.self_attn_norm(x) * (1 + scale_msa) + shift_msa
+        x_hidden = self.self_attn(x=x_hidden, pos_emb=pos_emb, window_size=self.window_size, kv_cache=kv_cache, return_kv=return_kv)
+        if return_kv: x_hidden, kv = x_hidden
+        x = x + x_hidden * gate_msa
+        x_hidden = self.cross_attn_norm(x)
+        x_hidden = self.cross_attn(x=x_hidden, y=y)
+        x = x + x_hidden
+        x_hidden = self.mlp_norm(x) * (1 + c_scale_msa) + c_shift_msa
+        x_hidden = self.mlp(x_hidden)
+        x = x + x_hidden * c_gate_msa
+        if return_kv:
+            return x, kv
+        else:
+            return x
+
+
+class DiffSynthMusicChannelProj(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, patch_size, bias=False, transposed=False):
+        super().__init__()
+        if transposed:
+            self.conv = torch.nn.ConvTranspose1d(in_channels=in_channels, out_channels=out_channels, kernel_size=patch_size, stride=patch_size, bias=bias)
+        else:
+            self.conv = torch.nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=patch_size, stride=patch_size, bias=bias)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)
+        x = self.conv(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class DiffSynthMusicDiTModel(torch.nn.Module):
+    def __init__(self, dim=2560, dim_mlp=9728, dim_condition=2048, num_blocks=32, num_heads_q=32, num_heads_kv=8, head_dim=128, window_size=128, patch_size=2):
+        super().__init__()
+        self.rotary_emb = Qwen3RotaryEmbedding(type('RopeConfig', (), {'head_dim': head_dim, 'max_position_embeddings': 32768, 'rope_theta': 1000000, 'rope_parameters': {'rope_type': 'default', 'rope_theta': 1000000}})())
+        self.x_emb = DiffSynthMusicChannelProj(in_channels=64*3, out_channels=dim, patch_size=patch_size, bias=True)
+        self.timestep_emb = DiffSynthMusicTimestepEmbedding(in_channels=256, time_embed_dim=dim)
+        self.condition_emb = torch.nn.Linear(dim_condition, dim, bias=True)
+        self.layers = torch.nn.ModuleList([
+            DiffSynthMusicDiTLayer(dim=dim, num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, dim_mlp=dim_mlp, window_size=window_size if block_id % 2 == 0 else None)
+            for block_id in range(num_blocks)
+        ])
+        self.norm_out = torch.nn.RMSNorm(dim, eps=1e-6)
+        self.proj_out = DiffSynthMusicChannelProj(in_channels=dim, out_channels=64, patch_size=patch_size, bias=True, transposed=True)
+        self.scale_shift_table = torch.nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+        self.placeholder_audio = torch.nn.Parameter(torch.zeros((1, 15000, 64)))
+
+    def forward_kv_cache(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        use_gradient_checkpointing: bool = False,
+        use_gradient_checkpointing_offload: bool = False,
+        **kwargs,
+    ):
+        x = torch.concat([self.placeholder_audio[:, :x.shape[1]], torch.ones_like(x), x], dim=-1)
+        x = self.x_emb(x)
+        y = self.condition_emb(y)
+        timestep = torch.zeros((1,), dtype=x.dtype, device=x.device)
+        final_timestep_emb, timestep_emb = self.timestep_emb(timestep)
+        pos_emb = self.rotary_emb(x, torch.arange(0, x.shape[1], device=x.device).unsqueeze(0))
+
+        kv_cache = {}
+        for block_id, block in enumerate(self.layers):
+            return_kv = block.window_size is None
+            x = gradient_checkpoint_forward(
+                block,
+                use_gradient_checkpointing,
+                use_gradient_checkpointing_offload,
+                x, y, pos_emb, timestep_emb,
+                return_kv=return_kv,
+                kv_cache=None,
+            )
+            if return_kv:
+                x, kv = x
+                kv_cache[f"{block_id}"] = kv
+
+        return kv_cache
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        timestep: torch.Tensor,
+        kv_cache = None,
+        use_gradient_checkpointing: bool = False,
+        use_gradient_checkpointing_offload: bool = False,
+        **kwargs,
+    ):
+        x = torch.concat([self.placeholder_audio[:, :x.shape[1]], torch.ones_like(x), x], dim=-1)
+        x = self.x_emb(x)
+        y = self.condition_emb(y)
+        final_timestep_emb, timestep_emb = self.timestep_emb(timestep)
+        pos_emb = self.rotary_emb(x, torch.arange(0, x.shape[1], device=x.device).unsqueeze(0))
+
+        for block_id, block in enumerate(self.layers):
+            x = gradient_checkpoint_forward(
+                block,
+                use_gradient_checkpointing,
+                use_gradient_checkpointing_offload,
+                x, y, pos_emb, timestep_emb,
+                return_kv=False,
+                kv_cache=None if kv_cache is None else kv_cache.get(f"{block_id}"),
+            )
+
+        shift, scale = (self.scale_shift_table.to(dtype=x.dtype, device=x.device) + final_timestep_emb.unsqueeze(1)).chunk(2, dim=1)
+        x = self.proj_out(self.norm_out(x) * (1 + scale) + shift)
+        return x