Delete files asr_0.wav assets/svqa_0.jpg assets/vqa_0.png chat_template.txt magvit_tokenizer/__init__.py magvit_tokenizer/common_modules.py magvit_tokenizer/misc.py magvit_tokenizer/modeling_magvitv2.py magvit_tokenizer/modeling_utils.py omni_diffusion/__init__.py omni_diffusion/constants.py omni_diffusion/data/__init__.py omni_diffusion/data/build.py omni_diffusion/data/data_collator.py omni_diffusion/data/dataset_base.py omni_diffusion/data/dataset_qwen2.py omni_diffusion/data/processor/__init__.py omni_diffusion/data/processor/audio_processor.py omni_diffusion/data/processor/image_processor.py omni_diffusion/data/utils.py omni_diffusion/models/__init__.py omni_diffusion/models/dream/__init__.py omni_diffusion/models/dream/config.yaml omni_diffusion/models/dream/config_dream_resume.json omni_diffusion/models/dream/configuration.json omni_diffusion/models/dream/configuration_dream.py omni_diffusion/models/dream/generation_utils.py omni_diffusion/models/dream/modeling_dream.py omni_diffusion/models/dream/modeling_sensevoice.py omni_diffusion/models/dream/resampler_projector.py omni_diffusion/models/dream/tokenization_dream.py omni_diffusion/models/dream/tokenizer_config.json omni_diffusion/models/magvit/common_modules.py omni_diffusion/models/magvit/misc.py omni_diffusion/models/magvit/modeling_magvitv2.py omni_diffusion/models/magvit/modeling_utils.py omni_diffusion/tokenizer.py omni_diffusion/tokenizer_magvitv2.py omni_diffusion/tokenizer_sensevoice_glm4voice.py s2i_0.wav svqa_0.jpg svqa_0.wav third_party/GLM-4-Voice/.gitignore third_party/GLM-4-Voice/LICENSE third_party/GLM-4-Voice/README.md third_party/GLM-4-Voice/README_en.md third_party/GLM-4-Voice/audio_process.py third_party/GLM-4-Voice/cosyvoice/__init__.py third_party/GLM-4-Voice/cosyvoice/bin/inference.py third_party/GLM-4-Voice/cosyvoice/bin/train.py third_party/GLM-4-Voice/cosyvoice/cli/__init__.py third_party/GLM-4-Voice/cosyvoice/cli/cosyvoice.py third_party/GLM-4-Voice/cosyvoice/cli/frontend.py third_party/GLM-4-Voice/cosyvoice/cli/model.py third_party/GLM-4-Voice/cosyvoice/dataset/__init__.py third_party/GLM-4-Voice/cosyvoice/dataset/dataset.py third_party/GLM-4-Voice/cosyvoice/dataset/processor.py third_party/GLM-4-Voice/cosyvoice/flow/decoder.py third_party/GLM-4-Voice/cosyvoice/flow/flow.py third_party/GLM-4-Voice/cosyvoice/flow/flow_gradtts.py third_party/GLM-4-Voice/cosyvoice/flow/flow_matching.py third_party/GLM-4-Voice/cosyvoice/flow/flow_matching_dit.py third_party/GLM-4-Voice/cosyvoice/flow/length_regulator.py third_party/GLM-4-Voice/cosyvoice/flow/stable/adp.py third_party/GLM-4-Voice/cosyvoice/flow/stable/blocks.py third_party/GLM-4-Voice/cosyvoice/flow/stable/dit.py third_party/GLM-4-Voice/cosyvoice/flow/stable/dit_v2.py third_party/GLM-4-Voice/cosyvoice/flow/stable/sampling.py third_party/GLM-4-Voice/cosyvoice/flow/stable/stable_diffusion.py third_party/GLM-4-Voice/cosyvoice/flow/stable/stable_diffusion_test.py third_party/GLM-4-Voice/cosyvoice/flow/stable/transformer.py third_party/GLM-4-Voice/cosyvoice/flow/stable/transformer_use_mask.py third_party/GLM-4-Voice/cosyvoice/hifigan/f0_predictor.py third_party/GLM-4-Voice/cosyvoice/hifigan/generator.py third_party/GLM-4-Voice/cosyvoice/llm/llm.py third_party/GLM-4-Voice/cosyvoice/transformer/__init__.py third_party/GLM-4-Voice/cosyvoice/transformer/activation.py third_party/GLM-4-Voice/cosyvoice/transformer/attention.py third_party/GLM-4-Voice/cosyvoice/transformer/convolution.py third_party/GLM-4-Voice/cosyvoice/transformer/decoder.py third_party/GLM-4-Voice/cosyvoice/transformer/decoder_layer.py third_party/GLM-4-Voice/cosyvoice/transformer/embedding.py third_party/GLM-4-Voice/cosyvoice/transformer/encoder.py third_party/GLM-4-Voice/cosyvoice/transformer/encoder_layer.py third_party/GLM-4-Voice/cosyvoice/transformer/label_smoothing_loss.py third_party/GLM-4-Voice/cosyvoice/transformer/positionwise_feed_forward.py third_party/GLM-4-Voice/cosyvoice/transformer/subsampling.py third_party/GLM-4-Voice/cosyvoice/utils/__init__.py third_party/GLM-4-Voice/cosyvoice/utils/block_mask_util.py third_party/GLM-4-Voice/cosyvoice/utils/class_utils.py third_party/GLM-4-Voice/cosyvoice/utils/common.py third_party/GLM-4-Voice/cosyvoice/utils/executor.py third_party/GLM-4-Voice/cosyvoice/utils/file_utils.py third_party/GLM-4-Voice/cosyvoice/utils/frontend_utils.py third_party/GLM-4-Voice/cosyvoice/utils/mask.py third_party/GLM-4-Voice/cosyvoice/utils/scheduler.py third_party/GLM-4-Voice/cosyvoice/utils/train_utils.py third_party/GLM-4-Voice/flow_inference.py third_party/GLM-4-Voice/model_server.py third_party/GLM-4-Voice/requirements.txt third_party/GLM-4-Voice/resources/architecture.jpeg third_party/GLM-4-Voice/resources/web_demo.png third_party/GLM-4-Voice/speech_tokenizer/__init__.py third_party/GLM-4-Voice/speech_tokenizer/configuration_whisper.py third_party/GLM-4-Voice/speech_tokenizer/generation_whisper.py third_party/GLM-4-Voice/speech_tokenizer/modeling_whisper.py third_party/GLM-4-Voice/speech_tokenizer/utils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/.env.example third_party/GLM-4-Voice/third_party/Matcha-TTS/.github/PULL_REQUEST_TEMPLATE.md third_party/GLM-4-Voice/third_party/Matcha-TTS/.github/codecov.yml third_party/GLM-4-Voice/third_party/Matcha-TTS/.github/dependabot.yml third_party/GLM-4-Voice/third_party/Matcha-TTS/.github/release-drafter.yml third_party/GLM-4-Voice/third_party/Matcha-TTS/.gitignore third_party/GLM-4-Voice/third_party/Matcha-TTS/.pre-commit-config.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/.project-root third_party/GLM-4-Voice/third_party/Matcha-TTS/.pylintrc third_party/GLM-4-Voice/third_party/Matcha-TTS/LICENSE third_party/GLM-4-Voice/third_party/Matcha-TTS/MANIFEST.in third_party/GLM-4-Voice/third_party/Matcha-TTS/Makefile third_party/GLM-4-Voice/third_party/Matcha-TTS/README.md third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/callbacks/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/callbacks/model_checkpoint.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/callbacks/model_summary.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/callbacks/none.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/callbacks/rich_progress_bar.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/data/hi-fi_en-US_female.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/data/ljspeech.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/data/vctk.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/debug/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/debug/fdr.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/debug/limit.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/debug/overfit.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/debug/profiler.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/eval.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/experiment/hifi_dataset_piper_phonemizer.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/experiment/ljspeech.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/experiment/ljspeech_from_durations.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/experiment/ljspeech_min_memory.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/experiment/multispeaker.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/extras/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/hparams_search/mnist_optuna.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/hydra/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/local/.gitkeep third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/aim.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/comet.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/csv.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/many_loggers.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/mlflow.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/neptune.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/tensorboard.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/logger/wandb.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/model/cfm/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/model/decoder/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/model/encoder/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/model/matcha.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/model/optimizer/adam.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/paths/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/train.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/cpu.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/ddp.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/ddp_sim.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/default.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/gpu.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/configs/trainer/mps.yaml third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/VERSION third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/app.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/cli.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/data/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/data/components/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/data/text_mel_datamodule.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/LICENSE third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/README.md third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/config.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/denoiser.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/env.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/meldataset.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/models.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/hifigan/xutils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/baselightningmodule.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/components/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/components/decoder.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/components/flow_matching.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/components/text_encoder.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/components/transformer.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/models/matcha_tts.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/onnx/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/onnx/export.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/onnx/infer.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/text/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/text/cleaners.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/text/numbers.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/text/symbols.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/train.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/audio.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/data/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/data/hificaptain.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/data/ljspeech.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/data/utils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/generate_data_statistics.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/get_durations_from_trained_model.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/instantiators.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/logging_utils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/model.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/monotonic_align/__init__.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/monotonic_align/core.pyx third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/monotonic_align/setup.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/pylogger.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/rich_utils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/matcha/utils/utils.py third_party/GLM-4-Voice/third_party/Matcha-TTS/notebooks/.gitkeep third_party/GLM-4-Voice/third_party/Matcha-TTS/pyproject.toml third_party/GLM-4-Voice/third_party/Matcha-TTS/requirements.txt third_party/GLM-4-Voice/third_party/Matcha-TTS/scripts/schedule.sh third_party/GLM-4-Voice/third_party/Matcha-TTS/setup.py third_party/GLM-4-Voice/third_party/Matcha-TTS/synthesis.ipynb third_party/GLM-4-Voice/web_demo.py vqa_0.png with huggingface_hub

asr_0.wav

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8cc5488d43dadd87e8fa93f8c2a418b6771b7a9201c87ef39fff7b10e5650803
-size 1108218

chat_template.txt

@@ -1,2 +0,0 @@
-
-{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n

magvit_tokenizer/common_modules.py

@@ -1,393 +0,0 @@
-"""
-Modified from https://github.com/CompVis/taming-transformers/blob/master/taming/modules/diffusionmodules/model.py#L34
-"""
-
-import math
-from typing import Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange, repeat
-from einops.layers.torch import Rearrange
-
-
-def nonlinearity(x):
-    # swish
-    return x * torch.sigmoid(x)
-
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(
-        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-    )
-
-
-class Upsample(nn.Module):
-    """
-    Spatial upsampling layer.
-    Upsamples by a factor of 2 using nearest neighbor interpolation,
-    optionally followed by a convolution.
-    """
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(
-                in_channels, in_channels, kernel_size=3, stride=1, padding=1
-            )
-
-    def forward(self, x):
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-
-class DepthToSpaceUpsample(nn.Module):
-    """
-    Upsamples spatially by moving channel depth to spatial dimensions (PixelShuffle).
-    Expands channels by 4, then rearranges (c * 2 * 2) -> (h * 2, w * 2).
-    """
-    def __init__(
-        self,
-        in_channels,
-    ):
-        super().__init__()
-        conv = nn.Conv2d(in_channels, in_channels * 4, 1)
-
-        self.net = nn.Sequential(
-            conv,
-            nn.SiLU(),
-            Rearrange("b (c p1 p2) h w -> b c (h p1) (w p2)", p1=2, p2=2),
-        )
-
-        self.init_conv_(conv)
-
-    def init_conv_(self, conv):
-        o, i, h, w = conv.weight.shape
-        conv_weight = torch.empty(o // 4, i, h, w)
-        nn.init.kaiming_uniform_(conv_weight)
-        conv_weight = repeat(conv_weight, "o ... -> (o 4) ...")
-
-        conv.weight.data.copy_(conv_weight)
-        nn.init.zeros_(conv.bias.data)
-
-    def forward(self, x):
-        out = self.net(x)
-        return out
-
-
-class Downsample(nn.Module):
-    """
-    Spatial downsampling layer.
-    Uses either a strided convolution (with padding) or average pooling.
-    """
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(
-                in_channels, in_channels, kernel_size=3, stride=2, padding=0
-            )
-
-    def forward(self, x):
-        if self.with_conv:
-            pad = (0, 1, 0, 1)
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-
-def unpack_time(t, batch):
-    _, c, w, h = t.size()
-    out = torch.reshape(t, [batch, -1, c, w, h])
-    out = rearrange(out, "b t c h w -> b c t h w")
-    return out
-
-
-def pack_time(t):
-    out = rearrange(t, "b c t h w -> b t c h w")
-    _, _, c, w, h = out.size()
-    return torch.reshape(out, [-1, c, w, h])
-
-
-class TimeDownsample2x(nn.Module):
-    """
-    Temporal downsampling layer.
-    Reduces the time dimension (T) by half using 1D convolution with causal padding.
-    """
-    def __init__(
-        self,
-        dim,
-        dim_out=None,
-        kernel_size=3,
-    ):
-        super().__init__()
-        if dim_out is None:
-            dim_out = dim
-        self.time_causal_padding = (kernel_size - 1, 0)
-        self.conv = nn.Conv1d(dim, dim_out, kernel_size, stride=2)
-
-    def forward(self, x):
-        x = rearrange(x, "b c t h w -> b h w c t")
-        b, h, w, c, t = x.size()
-        x = torch.reshape(x, [-1, c, t])
-
-        x = F.pad(x, self.time_causal_padding)
-        out = self.conv(x)
-
-        out = torch.reshape(out, [b, h, w, c, t])
-        out = rearrange(out, "b h w c t -> b c t h w")
-        out = rearrange(out, "b h w c t -> b c t h w")
-        return out
-
-
-class TimeUpsample2x(nn.Module):
-    """
-    Temporal upsampling layer.
-    Doubles the time dimension (T) using a PixelShuffle-like approach on the time axis.
-    """
-    def __init__(self, dim, dim_out=None):
-        super().__init__()
-        if dim_out is None:
-            dim_out = dim
-        conv = nn.Conv1d(dim, dim_out * 2, 1)
-
-        self.net = nn.Sequential(
-            nn.SiLU(), conv, Rearrange("b (c p) t -> b c (t p)", p=2)
-        )
-
-        self.init_conv_(conv)
-
-    def init_conv_(self, conv):
-        o, i, t = conv.weight.shape
-        conv_weight = torch.empty(o // 2, i, t)
-        nn.init.kaiming_uniform_(conv_weight)
-        conv_weight = repeat(conv_weight, "o ... -> (o 2) ...")
-
-        conv.weight.data.copy_(conv_weight)
-        nn.init.zeros_(conv.bias.data)
-
-    def forward(self, x):
-        x = rearrange(x, "b c t h w -> b h w c t")
-        b, h, w, c, t = x.size()
-        x = torch.reshape(x, [-1, c, t])
-
-        out = self.net(x)
-        out = out[:, :, 1:].contiguous()
-
-        out = torch.reshape(out, [b, h, w, c, t])
-        out = rearrange(out, "b h w c t -> b c t h w")
-        return out
-
-
-class AttnBlock(nn.Module):
-    """
-    Standard Self-Attention Block for spatial dimensions.
-    """
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.k = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.v = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.proj_out = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, h, w = q.shape
-        q = q.reshape(b, c, h * w)
-        q = q.permute(0, 2, 1)  # b,hw,c
-        k = k.reshape(b, c, h * w)  # b,c,hw
-        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b, c, h * w)
-        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b, c, h, w)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class TimeAttention(AttnBlock):
-    """
-    Applies the AttnBlock over the temporal dimension (T).
-    It treats spatial dimensions (H, W) as part of the batch.
-    """
-    def forward(self, x, *args, **kwargs):
-        x = rearrange(x, "b c t h w -> b h w t c")
-        b, h, w, t, c = x.size()
-        x = torch.reshape(x, (-1, t, c))
-
-        x = super().forward(x, *args, **kwargs)
-
-        x = torch.reshape(x, [b, h, w, t, c])
-        return rearrange(x, "b h w t c -> b c t h w")
-
-
-class Residual(nn.Module):
-    def __init__(self, fn: nn.Module):
-        super().__init__()
-        self.fn = fn
-
-    def forward(self, x, **kwargs):
-        return self.fn(x, **kwargs) + x
-
-
-def cast_tuple(t, length=1):
-    return t if isinstance(t, tuple) else ((t,) * length)
-
-
-class CausalConv3d(nn.Module):
-    """
-    3D Convolution with Causal Padding on the time axis.
-    Ensures that the convolution at time T only depends on frames <= T.
-    """
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int, int, int]],
-        pad_mode="constant",
-        **kwargs
-    ):
-        super().__init__()
-        kernel_size = cast_tuple(kernel_size, 3)
-
-        time_kernel_size, height_kernel_size, width_kernel_size = kernel_size
-
-        dilation = kwargs.pop("dilation", 1)
-        stride = kwargs.pop("stride", 1)
-
-        self.pad_mode = pad_mode
-        time_pad = dilation * (time_kernel_size - 1) + (1 - stride)
-        height_pad = height_kernel_size // 2
-        width_pad = width_kernel_size // 2
-
-        self.time_pad = time_pad
-        self.time_causal_padding = (
-            width_pad,
-            width_pad,
-            height_pad,
-            height_pad,
-            time_pad,
-            0,
-        )
-
-        stride = (stride, 1, 1)
-        dilation = (dilation, 1, 1)
-        self.conv = nn.Conv3d(
-            chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs
-        )
-
-    def forward(self, x):
-        pad_mode = self.pad_mode if self.time_pad < x.shape[2] else "constant"
-
-        x = F.pad(x, self.time_causal_padding, mode=pad_mode)
-        return self.conv(x)
-
-
-def ResnetBlockCausal3D(
-    dim, kernel_size: Union[int, Tuple[int, int, int]], pad_mode: str = "constant"
-):
-    net = nn.Sequential(
-        Normalize(dim),
-        nn.SiLU(),
-        CausalConv3d(dim, dim, kernel_size, pad_mode),
-        Normalize(dim),
-        nn.SiLU(),
-        CausalConv3d(dim, dim, kernel_size, pad_mode),
-    )
-    return Residual(net)
-
-
-class ResnetBlock(nn.Module):
-    """
-    Standard ResNet Block commonly used in Diffusion Models.
-    Includes support for Time Embeddings (temb) injection.
-    """
-    def __init__(
-        self,
-        *,
-        in_channels,
-        out_channels=None,
-        conv_shortcut=False,
-        dropout,
-        temb_channels=512
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(
-            in_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
-        else:
-            self.temb_proj = None
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(
-            out_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
-                )
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
-                )
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
-
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x + h

magvit_tokenizer/misc.py

@@ -1,55 +0,0 @@
-# copy from MMaDA
-
-from omegaconf import OmegaConf
-import torch
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    NewType,
-    Optional,
-    Sized,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-)
-try:
-    from typing import Literal
-except ImportError:
-    from typing_extensions import Literal
-
-# Tensor dtype
-# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
-from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
-
-# Config type
-from omegaconf import DictConfig
-
-# PyTorch Tensor type
-from torch import Tensor
-
-# Runtime type checking decorator
-from typeguard import typechecked as typechecker
-
-
-def broadcast(tensor, src=0):
-    if not _distributed_available():
-        return tensor
-    else:
-        torch.distributed.broadcast(tensor, src=src)
-        return tensor
-
-def _distributed_available():
-    return torch.distributed.is_available() and torch.distributed.is_initialized()
-
-def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
-    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
-    if '--local-rank' in cfg:
-        del cfg['--local-rank']
-    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
-    scfg = OmegaConf.structured(fields(**cfg))
-    return scfg

magvit_tokenizer/modeling_magvitv2.py

@@ -1,445 +0,0 @@
-"""
-Modified from https://github.com/Gen-Verse/MMaDA/blob/main/models/modeling_magvitv2.py
-"""
-
-from dataclasses import dataclass, field
-import numpy as np
-import torch
-import torch.nn as nn
-from .common_modules import *
-from .modeling_utils import ConfigMixin, ModelMixin, register_to_config
-from .misc import *
-import math
-
-class Updateable:
-    def do_update_step(
-            self, epoch: int, global_step: int, on_load_weights: bool = False
-    ):
-        for attr in self.__dir__():
-            if attr.startswith("_"):
-                continue
-            try:
-                module = getattr(self, attr)
-            except:
-                continue  # ignore attributes like property, which can't be retrived using getattr?
-            if isinstance(module, Updateable):
-                module.do_update_step(
-                    epoch, global_step, on_load_weights=on_load_weights
-                )
-        self.update_step(epoch, global_step, on_load_weights=on_load_weights)
-
-    def do_update_step_end(self, epoch: int, global_step: int):
-        for attr in self.__dir__():
-            if attr.startswith("_"):
-                continue
-            try:
-                module = getattr(self, attr)
-            except:
-                continue  # ignore attributes like property, which can't be retrived using getattr?
-            if isinstance(module, Updateable):
-                module.do_update_step_end(epoch, global_step)
-        self.update_step_end(epoch, global_step)
-
-    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
-        # override this method to implement custom update logic
-        # if on_load_weights is True, you should be careful doing things related to model evaluations,
-        # as the models and tensors are not guarenteed to be on the same device
-        pass
-
-    def update_step_end(self, epoch: int, global_step: int):
-        pass
-
-class VQGANEncoder(ModelMixin, ConfigMixin):
-    @dataclass
-    class Config:
-        ch: int = 128
-        ch_mult: List[int] = field(default_factory=lambda: [1, 2, 2, 4, 4])
-        num_res_blocks: List[int] = field(default_factory=lambda: [4, 3, 4, 3, 4])
-        attn_resolutions: List[int] = field(default_factory=lambda: [5])
-        dropout: float = 0.0
-        in_ch: int = 3
-        out_ch: int = 3
-        resolution: int = 256
-        z_channels: int = 13
-        double_z: bool = False
-
-    def __init__(self,
-                 ch: int = 128,
-                 ch_mult: List[int] = [1, 2, 2, 4, 4],
-                 num_res_blocks: List[int] = [4, 3, 4, 3, 4],
-                 attn_resolutions: List[int] = [5],
-                 dropout: float = 0.0,
-                 in_ch: int = 3,
-                 out_ch: int = 3,
-                 resolution: int = 256,
-                 z_channels: int = 13,
-                 double_z: bool = False):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_ch = in_ch
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(
-            self.in_ch, self.ch, kernel_size=3, stride=1, padding=1
-        )
-
-        curr_res = self.resolution
-        in_ch_mult = (1,) + tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = self.ch * in_ch_mult[i_level]
-            block_out = self.ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks[i_level]):
-                block.append(
-                    ResnetBlock(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        temb_channels=self.temb_ch,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                down.downsample = Downsample(block_in, True)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-
-
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(
-            block_in,
-            2 * z_channels if double_z else z_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-        )
-
-        self.quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
-        # for param in self.parameters():
-        #     broadcast(param, src=0)
-
-    def forward(self, x):
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks[i_level]):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions - 1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        h = self.quant_conv(h)
-        return h
-
-
-class LFQuantizer(nn.Module):
-    def __init__(self, num_codebook_entry: int = -1,
-                 codebook_dim: int = 13,
-                 beta: float = 0.25,
-                 entropy_multiplier: float = 0.1,
-                 commit_loss_multiplier: float = 0.1, ):
-        super().__init__()
-        self.codebook_size = 2 ** codebook_dim
-        print(
-            f"Look-up free quantizer with codebook size: {self.codebook_size}"
-        )
-        self.e_dim = codebook_dim
-        self.beta = beta
-
-        indices = torch.arange(self.codebook_size)
-
-        binary = (
-                         indices.unsqueeze(1)
-                         >> torch.arange(codebook_dim - 1, -1, -1, dtype=torch.long)
-                 ) & 1
-
-        embedding = binary.float() * 2 - 1
-        self.register_buffer("embedding", embedding)
-        self.register_buffer(
-            "power_vals", 2 ** torch.arange(codebook_dim - 1, -1, -1)
-        )
-        self.commit_loss_multiplier = commit_loss_multiplier
-        self.entropy_multiplier = entropy_multiplier
-
-    def get_indices(self, z_q):
-        return (
-            (self.power_vals.reshape(1, -1, 1, 1) * (z_q > 0).float())
-            .sum(1, keepdim=True)
-            .long()
-        )
-
-    def get_codebook_entry(self, indices, shape=None):
-        if shape is None:
-            h, w = int(math.sqrt(indices.shape[-1])), int(math.sqrt(indices.shape[-1]))
-        else:
-            h, w = shape
-        b, _ = indices.shape
-        indices = indices.reshape(-1)
-        z_q = self.embedding[indices]
-        z_q = z_q.view(b, h, w, -1)
-
-        # reshape back to match original input shape
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q
-
-    def forward(self, z, get_code=False):
-        """
-        Inputs the output of the encoder network z and maps it to a discrete
-        one-hot vector that is the index of the closest embedding vector e_j
-        z (continuous) -> z_q (discrete)
-        z.shape = (batch, channel, height, width)
-        quantization pipeline:
-            1. get encoder input (B,C,H,W)
-            2. flatten input to (B*H*W,C)
-        """
-        if get_code:
-            return self.get_codebook_entry(z)
-
-        # reshape z -> (batch, height, width, channel) and flatten
-        z = z.permute(0, 2, 3, 1).contiguous()
-        z_flattened = z.view(-1, self.e_dim)
-        ge_zero = (z_flattened > 0).float()
-        ones = torch.ones_like(z_flattened)
-        z_q = ones * ge_zero + -ones * (1 - ge_zero)
-
-        # preserve gradients
-        z_q = z_flattened + (z_q - z_flattened).detach()
-
-        # compute entropy loss
-        CatDist = torch.distributions.categorical.Categorical
-        logit = torch.stack(
-            [
-                -(z_flattened - torch.ones_like(z_q)).pow(2),
-                -(z_flattened - torch.ones_like(z_q) * -1).pow(2),
-            ],
-            dim=-1,
-        )
-        cat_dist = CatDist(logits=logit)
-        entropy = cat_dist.entropy().mean()
-        mean_prob = cat_dist.probs.mean(0)
-        mean_entropy = CatDist(probs=mean_prob).entropy().mean()
-
-        # compute loss for embedding
-        commit_loss = torch.mean(
-            (z_q.detach() - z_flattened) ** 2
-        ) + self.beta * torch.mean((z_q - z_flattened.detach()) ** 2)
-
-        # reshape back to match original input shape
-        z_q = z_q.view(z.shape)
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return {
-            "z": z_q,
-            "quantizer_loss": commit_loss * self.commit_loss_multiplier,
-            "entropy_loss": (entropy - mean_entropy) * self.entropy_multiplier,
-            "indices": self.get_indices(z_q),
-        }
-
-
-class VQGANDecoder(ModelMixin, ConfigMixin):
-    def __init__(self, ch: int = 128,
-                 ch_mult: List[int] = [1, 1, 2, 2, 4],
-                 num_res_blocks: List[int] = [4, 4, 3, 4, 3],
-                 attn_resolutions: List[int] = [5],
-                 dropout: float = 0.0,
-                 in_ch: int = 3,
-                 out_ch: int = 3,
-                 resolution: int = 256,
-                 z_channels: int = 13,
-                 double_z: bool = False):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_ch = in_ch
-        self.give_pre_end = False
-
-        self.z_channels = z_channels
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,) + tuple(ch_mult)
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        curr_res = self.resolution // 2 ** (self.num_resolutions - 1)
-        self.z_shape = (1, z_channels, curr_res, curr_res)
-        print(
-            "Working with z of shape {} = {} dimensions.".format(
-                self.z_shape, np.prod(self.z_shape)
-            )
-        )
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv2d(
-            z_channels, block_in, kernel_size=3, stride=1, padding=1
-        )
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks[i_level]):
-                block.append(
-                    ResnetBlock(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        temb_channels=self.temb_ch,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, True)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(
-            block_in, out_ch, kernel_size=3, stride=1, padding=1
-        )
-        self.post_quant_conv = torch.nn.Conv2d(
-            z_channels, z_channels, 1
-        )
-
-
-    def forward(self, z):
-        # assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-        # timestep embedding
-        temb = None
-        output = dict()
-        z = self.post_quant_conv(z)
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks[i_level]):
-                h = self.up[i_level].block[i_block](h, temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        output["output"] = h
-        if self.give_pre_end:
-            return output
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        output["output"] = h
-        return output
-
-
-class MAGVITv2(ModelMixin, ConfigMixin):
-    @register_to_config
-    def __init__(
-            self,
-    ):
-        super().__init__()
-
-        self.encoder = VQGANEncoder()
-        self.decoder = VQGANDecoder()
-        self.quantize = LFQuantizer()
-
-    def forward(self, pixel_values, return_loss=False):
-        pass
-
-    def encode(self, pixel_values, return_loss=False):
-        hidden_states = self.encoder(pixel_values)
-        quantized_states = self.quantize(hidden_states)['z']
-        codebook_indices = self.quantize.get_indices(quantized_states).reshape(pixel_values.shape[0], -1)
-        output = (quantized_states, codebook_indices)
-        return output
-
-    def get_code(self, pixel_values):
-        hidden_states = self.encoder(pixel_values)
-        codebook_indices = self.quantize.get_indices(self.quantize(hidden_states)['z'])
-        codebook_indices = codebook_indices.reshape(pixel_values.shape[0], -1)
-
-        return codebook_indices
-
-    def decode_code(self, codebook_indices, shape=None):
-        z_q = self.quantize.get_codebook_entry(codebook_indices, shape=shape)
-
-        reconstructed_pixel_values = self.decoder(z_q)["output"]
-        return reconstructed_pixel_values
-
-
-if __name__ == '__main__':
-    encoder = VQGANEncoder()
-    import ipdb
-    ipdb.set_trace()
-    print()

magvit_tokenizer/modeling_utils.py

@@ -1,1184 +0,0 @@
-"""
-Modified from https://github.com/Gen-Verse/MMaDA/blob/main/models/modeling_utils.py
-"""
-
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team.
-# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-import itertools
-import json
-import os
-import re
-from collections import OrderedDict
-from functools import partial
-from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
-
-import safetensors
-import torch
-from huggingface_hub import create_repo, split_torch_state_dict_into_shards
-from huggingface_hub.utils import validate_hf_hub_args
-from torch import Tensor, nn
-
-from diffusers import __version__
-from diffusers.utils import (
-    FLAX_WEIGHTS_NAME,
-    SAFE_WEIGHTS_INDEX_NAME,
-    WEIGHTS_INDEX_NAME,
-    _add_variant,
-    _get_checkpoint_shard_files,
-    _get_model_file,
-    deprecate,
-    is_accelerate_available,
-    is_torch_version,
-    logging,
-)
-
-CONFIG_NAME = "config.json"
-WEIGHTS_NAME = "pytorch_model.bin"
-SAFETENSORS_WEIGHTS_NAME = "pytorch_model.safetensors"
-HUGGINGFACE_CO_RESOLVE_ENDPOINT = "https://huggingface.co"
-
-from diffusers.utils.hub_utils import (
-    PushToHubMixin,
-    load_or_create_model_card,
-    populate_model_card,
-)
-from diffusers.models.model_loading_utils import (
-    _determine_device_map,
-    _fetch_index_file,
-    _load_state_dict_into_model,
-    load_model_dict_into_meta,
-)
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-import diffusers
-
-logger = logging.get_logger(__name__)
-
-_REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}")
-
-
-if is_torch_version(">=", "1.9.0"):
-    _LOW_CPU_MEM_USAGE_DEFAULT = True
-else:
-    _LOW_CPU_MEM_USAGE_DEFAULT = False
-
-
-if is_accelerate_available():
-    import accelerate
-
-
-def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
-    try:
-        parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
-        return next(parameters_and_buffers).device
-    except StopIteration:
-        # For torch.nn.DataParallel compatibility in PyTorch 1.5
-
-        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
-            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
-            return tuples
-
-        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
-        first_tuple = next(gen)
-        return first_tuple[1].device
-
-
-def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
-    try:
-        params = tuple(parameter.parameters())
-        if len(params) > 0:
-            return params[0].dtype
-
-        buffers = tuple(parameter.buffers())
-        if len(buffers) > 0:
-            return buffers[0].dtype
-
-    except StopIteration:
-        # For torch.nn.DataParallel compatibility in PyTorch 1.5
-
-        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
-            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
-            return tuples
-
-        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
-        first_tuple = next(gen)
-        return first_tuple[1].dtype
-
-
-class ModelMixin(torch.nn.Module, PushToHubMixin):
-    r"""
-    Base class for all models.
-
-    [`ModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and
-    saving models.
-
-        - **config_name** ([`str`]) -- Filename to save a model to when calling [`~models.ModelMixin.save_pretrained`].
-    """
-
-    config_name = CONFIG_NAME
-    _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
-    _supports_gradient_checkpointing = False
-    _keys_to_ignore_on_load_unexpected = None
-    _no_split_modules = None
-
-    def __init__(self):
-        super().__init__()
-
-    def __getattr__(self, name: str) -> Any:
-        """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
-        config attributes directly. See https://github.com/huggingface/diffusers/pull/3129 We need to overwrite
-        __getattr__ here in addition so that we don't trigger `torch.nn.Module`'s __getattr__':
-        https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
-        """
-
-        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
-        is_attribute = name in self.__dict__
-
-        if is_in_config and not is_attribute:
-            deprecation_message = f"Accessing config attribute `{name}` directly via \
-                                '{type(self).__name__}' object attribute is deprecated. \
-                                Please access '{name}' over '{type(self).__name__}'s config object instead, \
-                                e.g. 'unet.config.{name}'."
-            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)
-            return self._internal_dict[name]
-
-        # call PyTorch's https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
-        return super().__getattr__(name)
-
-    @property
-    def is_gradient_checkpointing(self) -> bool:
-        """
-        Whether gradient checkpointing is activated for this model or not.
-        """
-        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
-
-    def enable_gradient_checkpointing(self) -> None:
-        """
-        Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
-        *checkpoint activations* in other frameworks).
-        """
-        if not self._supports_gradient_checkpointing:
-            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
-        self.apply(partial(self._set_gradient_checkpointing, value=True))
-
-    def disable_gradient_checkpointing(self) -> None:
-        """
-        Deactivates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
-        *checkpoint activations* in other frameworks).
-        """
-        if self._supports_gradient_checkpointing:
-            self.apply(partial(self._set_gradient_checkpointing, value=False))
-
-    # def set_use_npu_flash_attention(self, valid: bool) -> None:
-    #     r"""
-    #     Set the switch for the npu flash attention.
-    #     """
-
-    #     def fn_recursive_set_npu_flash_attention(module: torch.nn.Module):
-    #         if hasattr(module, "set_use_npu_flash_attention"):
-    #             module.set_use_npu_flash_attention(valid)
-
-    #         for child in module.children():
-    #             fn_recursive_set_npu_flash_attention(child)
-
-    #     for module in self.children():
-    #         if isinstance(module, torch.nn.Module):
-    #             fn_recursive_set_npu_flash_attention(module)
-
-    # def enable_npu_flash_attention(self) -> None:
-    #     r"""
-    #     Enable npu flash attention from torch_npu
-
-    #     """
-    #     self.set_use_npu_flash_attention(True)
-
-    # def disable_npu_flash_attention(self) -> None:
-    #     r"""
-    #     disable npu flash attention from torch_npu
-
-    #     """
-    #     self.set_use_npu_flash_attention(False)
-
-    # def set_use_memory_efficient_attention_xformers(
-    #     self, valid: bool, attention_op: Optional[Callable] = None
-    # ) -> None:
-    #     # Recursively walk through all the children.
-    #     # Any children which exposes the set_use_memory_efficient_attention_xformers method
-    #     # gets the message
-    #     def fn_recursive_set_mem_eff(module: torch.nn.Module):
-    #         if hasattr(module, "set_use_memory_efficient_attention_xformers"):
-    #             module.set_use_memory_efficient_attention_xformers(valid, attention_op)
-
-    #         for child in module.children():
-    #             fn_recursive_set_mem_eff(child)
-
-    #     for module in self.children():
-    #         if isinstance(module, torch.nn.Module):
-    #             fn_recursive_set_mem_eff(module)
-
-    # def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None) -> None:
-    #     r"""
-    #     Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
-
-    #     When this option is enabled, you should observe lower GPU memory usage and a potential speed up during
-    #     inference. Speed up during training is not guaranteed.
-
-    #     <Tip warning={true}>
-
-    #     ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes
-    #     precedent.
-
-    #     </Tip>
-
-    #     Parameters:
-    #         attention_op (`Callable`, *optional*):
-    #             Override the default `None` operator for use as `op` argument to the
-    #             [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#x
-    #             formers.ops.memory_efficient_attention)
-    #             function of xFormers.
-
-    #     Examples:
-
-    #     ```py
-    #     >>> import torch
-    #     >>> from diffusers import UNet2DConditionModel
-    #     >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
-
-    #     >>> model = UNet2DConditionModel.from_pretrained(
-    #     ...     "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
-    #     ... )
-    #     >>> model = model.to("cuda")
-    #     >>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
-    #     ```
-    #     """
-    #     self.set_use_memory_efficient_attention_xformers(True, attention_op)
-
-    # def disable_xformers_memory_efficient_attention(self) -> None:
-    #     r"""
-    #     Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
-    #     """
-    #     self.set_use_memory_efficient_attention_xformers(False)
-
-    def save_pretrained(
-        self,
-        save_directory: Union[str, os.PathLike],
-        is_main_process: bool = True,
-        save_function: Optional[Callable] = None,
-        safe_serialization: bool = True,
-        variant: Optional[str] = None,
-        max_shard_size: Union[int, str] = "10GB",
-        push_to_hub: bool = False,
-        **kwargs,
-    ):
-        """
-        Save a model and its configuration file to a directory so that it can be reloaded using the
-        [`~models.ModelMixin.from_pretrained`] class method.
-
-        Arguments:
-            save_directory (`str` or `os.PathLike`):
-                Directory to save a model and its configuration file to. Will be created if it doesn't exist.
-            is_main_process (`bool`, *optional*, defaults to `True`):
-                Whether the process calling this is the main process or not. Useful during distributed training and you
-                need to call this function on all processes. In this case, set `is_main_process=True` only on the main
-                process to avoid race conditions.
-            save_function (`Callable`):
-                The function to use to save the state dictionary. Useful during distributed training when you need to
-                replace `torch.save` with another method. Can be configured with the environment variable
-                `DIFFUSERS_SAVE_MODE`.
-            safe_serialization (`bool`, *optional*, defaults to `True`):
-                Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
-            variant (`str`, *optional*):
-                If specified, weights are saved in the format `pytorch_model.<variant>.bin`.
-            max_shard_size (`int` or `str`, defaults to `"10GB"`):
-                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
-                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5GB"`).
-                If expressed as an integer, the unit is bytes. Note that this limit will be decreased after a certain
-                period of time (starting from Oct 2024) to allow users to upgrade to the latest version of `diffusers`.
-                This is to establish a common default size for this argument across different libraries in the Hugging
-                Face ecosystem (`transformers`, and `accelerate`, for example).
-            push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
-                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
-                namespace).
-            kwargs (`Dict[str, Any]`, *optional*):
-                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
-        """
-        if os.path.isfile(save_directory):
-            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
-            return
-
-        weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
-        weights_name = _add_variant(weights_name, variant)
-        weight_name_split = weights_name.split(".")
-        if len(weight_name_split) in [2, 3]:
-            weights_name_pattern = weight_name_split[0] + "{suffix}." + ".".join(weight_name_split[1:])
-        else:
-            raise ValueError(f"Invalid {weights_name} provided.")
-
-        os.makedirs(save_directory, exist_ok=True)
-
-        if push_to_hub:
-            commit_message = kwargs.pop("commit_message", None)
-            private = kwargs.pop("private", False)
-            create_pr = kwargs.pop("create_pr", False)
-            token = kwargs.pop("token", None)
-            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
-            repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
-
-        # Only save the model itself if we are using distributed training
-        model_to_save = self
-
-        # Attach architecture to the config
-        # Save the config
-        if is_main_process:
-            model_to_save.save_config(save_directory)
-
-        # Save the model
-        state_dict = model_to_save.state_dict()
-
-        # Save the model
-        state_dict_split = split_torch_state_dict_into_shards(
-            state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern
-        )
-
-        # Clean the folder from a previous save
-        if is_main_process:
-            for filename in os.listdir(save_directory):
-                if filename in state_dict_split.filename_to_tensors.keys():
-                    continue
-                full_filename = os.path.join(save_directory, filename)
-                if not os.path.isfile(full_filename):
-                    continue
-                weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "")
-                weights_without_ext = weights_without_ext.replace("{suffix}", "")
-                filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "")
-                # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
-                if (
-                    filename.startswith(weights_without_ext)
-                    and _REGEX_SHARD.fullmatch(filename_without_ext) is not None
-                ):
-                    os.remove(full_filename)
-
-        for filename, tensors in state_dict_split.filename_to_tensors.items():
-            shard = {tensor: state_dict[tensor] for tensor in tensors}
-            filepath = os.path.join(save_directory, filename)
-            if safe_serialization:
-                # At some point we will need to deal better with save_function (used for TPU and other distributed
-                # joyfulness), but for now this enough.
-                safetensors.torch.save_file(shard, filepath, metadata={"format": "pt"})
-            else:
-                torch.save(shard, filepath)
-
-        if state_dict_split.is_sharded:
-            index = {
-                "metadata": state_dict_split.metadata,
-                "weight_map": state_dict_split.tensor_to_filename,
-            }
-            save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
-            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-            # Save the index as well
-            with open(save_index_file, "w", encoding="utf-8") as f:
-                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                f.write(content)
-            logger.info(
-                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each "
-                f"parameters has been saved in the index located at {save_index_file}."
-            )
-        else:
-            path_to_weights = os.path.join(save_directory, weights_name)
-            logger.info(f"Model weights saved in {path_to_weights}")
-
-        if push_to_hub:
-            # Create a new empty model card and eventually tag it
-            model_card = load_or_create_model_card(repo_id, token=token)
-            model_card = populate_model_card(model_card)
-            model_card.save(Path(save_directory, "README.md").as_posix())
-
-            self._upload_folder(
-                save_directory,
-                repo_id,
-                token=token,
-                commit_message=commit_message,
-                create_pr=create_pr,
-            )
-
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        r"""
-        Instantiate a pretrained PyTorch model from a pretrained model configuration.
-
-        The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To
-        train the model, set it back in training mode with `model.train()`.
-
-        Parameters:
-            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
-                Can be either:
-
-                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
-                      the Hub.
-                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
-                      with [`~ModelMixin.save_pretrained`].
-
-            cache_dir (`Union[str, os.PathLike]`, *optional*):
-                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
-                is not used.
-            torch_dtype (`str` or `torch.dtype`, *optional*):
-                Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
-                dtype is automatically derived from the model's weights.
-            force_download (`bool`, *optional*, defaults to `False`):
-                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
-                cached versions if they exist.
-            proxies (`Dict[str, str]`, *optional*):
-                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
-                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
-            output_loading_info (`bool`, *optional*, defaults to `False`):
-                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
-            local_files_only(`bool`, *optional*, defaults to `False`):
-                Whether to only load local model weights and configuration files or not. If set to `True`, the model
-                won't be downloaded from the Hub.
-            token (`str` or *bool*, *optional*):
-                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
-                `diffusers-cli login` (stored in `~/.huggingface`) is used.
-            revision (`str`, *optional*, defaults to `"main"`):
-                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
-                allowed by Git.
-            from_flax (`bool`, *optional*, defaults to `False`):
-                Load the model weights from a Flax checkpoint save file.
-            subfolder (`str`, *optional*, defaults to `""`):
-                The subfolder location of a model file within a larger model repository on the Hub or locally.
-            mirror (`str`, *optional*):
-                Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
-                guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
-                information.
-            device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
-                A map that specifies where each submodule should go. It doesn't need to be defined for each
-                parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
-                same device. Defaults to `None`, meaning that the model will be loaded on CPU.
-
-                Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For
-                more information about each option see [designing a device
-                map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
-            max_memory (`Dict`, *optional*):
-                A dictionary device identifier for the maximum memory. Will default to the maximum memory available for
-                each GPU and the available CPU RAM if unset.
-            offload_folder (`str` or `os.PathLike`, *optional*):
-                The path to offload weights if `device_map` contains the value `"disk"`.
-            offload_state_dict (`bool`, *optional*):
-                If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if
-                the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True`
-                when there is some disk offload.
-            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
-                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
-                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
-                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
-                argument to `True` will raise an error.
-            variant (`str`, *optional*):
-                Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when
-                loading `from_flax`.
-            use_safetensors (`bool`, *optional*, defaults to `None`):
-                If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the
-                `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors`
-                weights. If set to `False`, `safetensors` weights are not loaded.
-
-        <Tip>
-
-        To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with
-        `huggingface-cli login`. You can also activate the special
-        ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a
-        firewalled environment.
-
-        </Tip>
-
-        Example:
-
-        ```py
-        from diffusers import UNet2DConditionModel
-
-        unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
-        ```
-
-        If you get the error message below, you need to finetune the weights for your downstream task:
-
-        ```bash
-        Some weights of UNet2DConditionModel were not initialized from the model checkpoint at 
-        runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match:
-        - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) 
-            in the model instantiated
-        You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
-        ```
-        """
-        cache_dir = kwargs.pop("cache_dir", None)
-        ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
-        force_download = kwargs.pop("force_download", False)
-        from_flax = kwargs.pop("from_flax", False)
-        proxies = kwargs.pop("proxies", None)
-        output_loading_info = kwargs.pop("output_loading_info", False)
-        local_files_only = kwargs.pop("local_files_only", None)
-        token = kwargs.pop("token", None)
-        revision = kwargs.pop("revision", None)
-        torch_dtype = kwargs.pop("torch_dtype", None)
-        subfolder = kwargs.pop("subfolder", None)
-        device_map = kwargs.pop("device_map", None)
-        max_memory = kwargs.pop("max_memory", None)
-        offload_folder = kwargs.pop("offload_folder", None)
-        offload_state_dict = kwargs.pop("offload_state_dict", False)
-        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
-        variant = kwargs.pop("variant", None)
-        use_safetensors = kwargs.pop("use_safetensors", None)
-
-        allow_pickle = False
-        if use_safetensors is None:
-            use_safetensors = True
-            allow_pickle = True
-
-        def _check_input(low_cpu_mem_usage, device_map):
-            if low_cpu_mem_usage and not is_accelerate_available():
-                low_cpu_mem_usage = False
-                logger.warning(
-                    "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
-                    " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
-                    " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
-                    " install accelerate\n```\n."
-                )
-
-            if device_map is not None and not is_accelerate_available():
-                raise NotImplementedError(
-                    "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set"
-                    " `device_map=None`. You can install accelerate with `pip install accelerate`."
-                )
-
-            # Check if we can handle device_map and dispatching the weights
-            if device_map is not None and not is_torch_version(">=", "1.9.0"):
-                raise NotImplementedError(
-                    "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
-                    " `device_map=None`."
-                )
-
-            if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
-                raise NotImplementedError(
-                    "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or "
-                    "set `low_cpu_mem_usage=False`."
-                )
-
-            if low_cpu_mem_usage is False and device_map is not None:
-                raise ValueError(
-                    f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and"
-                    " dispatching. Please make sure to set `low_cpu_mem_usage=True`."
-                )
-        _check_input(low_cpu_mem_usage, device_map)
-
-        # change device_map into a map if we passed an int, a str or a torch.device
-        if isinstance(device_map, torch.device):
-            device_map = {"": device_map}
-        elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
-            try:
-                device_map = {"": torch.device(device_map)}
-            except RuntimeError:
-                raise ValueError(
-                    "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or "
-                    f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}."
-                )
-        elif isinstance(device_map, int):
-            if device_map < 0:
-                raise ValueError(
-                    "You can't pass device_map as a negative int. " \
-                    "If you want to put the model on the cpu, pass device_map = 'cpu' "
-                )
-            else:
-                device_map = {"": device_map}
-
-        if device_map is not None:
-            if low_cpu_mem_usage is None:
-                low_cpu_mem_usage = True
-            elif not low_cpu_mem_usage:
-                raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`")
-
-        if low_cpu_mem_usage:
-            if device_map is not None and not is_torch_version(">=", "1.10"):
-                # The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info.
-                raise ValueError("`low_cpu_mem_usage` and `device_map` require PyTorch >= 1.10.")
-
-        # Load config if we don't provide a configuration
-        config_path = pretrained_model_name_or_path
-
-        user_agent = {
-            "diffusers": __version__,
-            "file_type": "model",
-            "framework": "pytorch",
-        }
-
-        # load config
-        def _load_config(cls, config_path, **kwargs):
-            config, unused_kwargs, commit_hash = cls.load_config(
-                config_path,
-                cache_dir=cache_dir,
-                return_unused_kwargs=True,
-                return_commit_hash=True,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                subfolder=subfolder,
-                user_agent=user_agent,
-                **kwargs,
-            )
-
-            # Determine if we're loading from a directory of sharded checkpoints.
-            is_sharded = False
-            index_file = None
-            is_local = os.path.isdir(pretrained_model_name_or_path)
-            index_file = _fetch_index_file(
-                is_local=is_local,
-                pretrained_model_name_or_path=pretrained_model_name_or_path,
-                subfolder=subfolder or "",
-                use_safetensors=use_safetensors,
-                cache_dir=cache_dir,
-                variant=variant,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                user_agent=user_agent,
-                commit_hash=commit_hash,
-            )
-            if index_file is not None and index_file.is_file():
-                is_sharded = True
-
-            if is_sharded and from_flax:
-                raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.")
-            return config, unused_kwargs, commit_hash, is_sharded, index_file
-        
-        config, unused_kwargs, commit_hash, is_sharded, index_file = _load_config(cls, config_path, **kwargs)
-
-        # load model
-        model_file = None
-        if from_flax:
-            model_file = _get_model_file(
-                pretrained_model_name_or_path,
-                weights_name=FLAX_WEIGHTS_NAME,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                subfolder=subfolder,
-                user_agent=user_agent,
-                commit_hash=commit_hash,
-            )
-            model = cls.from_config(config, **unused_kwargs)
-
-            # Convert the weights
-            from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model
-
-            model = load_flax_checkpoint_in_pytorch_model(model, model_file)
-        else:
-            if is_sharded:
-                sharded_ckpt_cached_folder, sharded_metadata = _get_checkpoint_shard_files(
-                    pretrained_model_name_or_path,
-                    index_file,
-                    cache_dir=cache_dir,
-                    proxies=proxies,
-                    local_files_only=local_files_only,
-                    token=token,
-                    user_agent=user_agent,
-                    revision=revision,
-                    subfolder=subfolder or "",
-                )
-
-            elif use_safetensors and not is_sharded:
-                try:
-                    model_file = _get_model_file(
-                        pretrained_model_name_or_path,
-                        weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
-                        cache_dir=cache_dir,
-                        force_download=force_download,
-                        proxies=proxies,
-                        local_files_only=local_files_only,
-                        token=token,
-                        revision=revision,
-                        subfolder=subfolder,
-                        user_agent=user_agent,
-                        commit_hash=commit_hash,
-                    )
-
-                except IOError as e:
-                    logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
-                    if not allow_pickle:
-                        raise
-                    logger.warning(
-                        "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
-                    )
-
-            if model_file is None and not is_sharded:
-                model_file = _get_model_file(
-                    pretrained_model_name_or_path,
-                    weights_name=_add_variant(WEIGHTS_NAME, variant),
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    proxies=proxies,
-                    local_files_only=local_files_only,
-                    token=token,
-                    revision=revision,
-                    subfolder=subfolder,
-                    user_agent=user_agent,
-                    commit_hash=commit_hash,
-                )
-
-            if low_cpu_mem_usage:
-                # Instantiate model with empty weights
-                def load_model_in_low_cpu_mem_usage(config, device_map, unused_kwargs, commit_hash, is_sharded, 
-                    index_file):
-                    with accelerate.init_empty_weights():
-                        model = cls.from_config(config, **unused_kwargs)
-
-                    # if device_map is None, load the state dict and move the params from meta device to the cpu
-                    if device_map is None and not is_sharded:
-                        param_device = "cpu"
-                        state_dict = diffusers.models.model_loading_utils.load_state_dict(model_file, variant=variant)
-                        model._convert_deprecated_attn(state_dict)
-                        # move the params from meta device to cpu
-                        missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
-                        if len(missing_keys) > 0:
-                            raise ValueError(
-                                f"Cannot load {cls} from {pretrained_model_name_or_path} because the following keys are"
-                                f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
-                                " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
-                                " those weights or else make sure your checkpoint file is correct."
-                            )
-
-                        unexpected_keys = load_model_dict_into_meta(
-                            model,
-                            state_dict,
-                            device=param_device,
-                            dtype=torch_dtype,
-                            model_name_or_path=pretrained_model_name_or_path,
-                        )
-
-                        if cls._keys_to_ignore_on_load_unexpected is not None:
-                            for pat in cls._keys_to_ignore_on_load_unexpected:
-                                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
-
-                        if len(unexpected_keys) > 0:
-                            logger.warning(
-                                f"Some weights of the model checkpoint were not used when initializing "
-                                f" {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
-                            )
-
-                    else:  # else let accelerate handle loading and dispatching.
-                        # Load weights and dispatch according to the device_map
-                        # by default the device_map is None and the weights are loaded on the CPU
-                        force_hook = True
-                        device_map = _determine_device_map(model, device_map, max_memory, torch_dtype)
-                        if device_map is None and is_sharded:
-                            # we load the parameters on the cpu
-                            device_map = {"": "cpu"}
-                            force_hook = False
-                        try:
-                            accelerate.load_checkpoint_and_dispatch(
-                                model,
-                                model_file if not is_sharded else index_file,
-                                device_map,
-                                max_memory=max_memory,
-                                offload_folder=offload_folder,
-                                offload_state_dict=offload_state_dict,
-                                dtype=torch_dtype,
-                                force_hooks=force_hook,
-                                strict=True,
-                            )
-                        except AttributeError as e:
-                            # When using accelerate loading, we do not have the ability to load the state
-                            # dict and rename the weight names manually. Additionally, accelerate skips
-                            # torch loading conventions and directly writes into `module.{_buffers, _parameters}`
-                            # (which look like they should be private variables?), so we can't use the standard hooks
-                            # to rename parameters on load. We need to mimic the original weight names so the correct
-                            # attributes are available. After we have loaded the weights, we convert the deprecated
-                            # names to the new non-deprecated names. Then we _greatly encourage_ the user to convert
-                            # the weights so we don't have to do this again.
-
-                            if "'Attention' object has no attribute" in str(e):
-                                logger.warning(
-                                    f"Taking `{str(e)}` while using `accelerate.load_checkpoint_and_dispatch` to mean "
-                                    f"{pretrained_model_name_or_path} was saved with deprecated attention block weight "
-                                    f"names. We will load it with the deprecated attention block names and convert them"
-                                    f" on the fly to the new attention block format. Please re-save the model after "
-                                    f"this conversion so we don't have to do the on the fly renaming in the future. "
-                                    f"If the model is from a hub checkpoint, please also re-upload it or open a PR on"
-                                    f" the original repository."
-                                )
-                                model._temp_convert_self_to_deprecated_attention_blocks()
-                                accelerate.load_checkpoint_and_dispatch(
-                                    model,
-                                    model_file if not is_sharded else index_file,
-                                    device_map,
-                                    max_memory=max_memory,
-                                    offload_folder=offload_folder,
-                                    offload_state_dict=offload_state_dict,
-                                    dtype=torch_dtype,
-                                    force_hooks=force_hook,
-                                    strict=True,
-                                )
-                                model._undo_temp_convert_self_to_deprecated_attention_blocks()
-                            else:
-                                raise e
-
-                    loading_info = {
-                        "missing_keys": [],
-                        "unexpected_keys": [],
-                        "mismatched_keys": [],
-                        "error_msgs": [],
-                    }
-                    return model, loading_info
-                model, loading_info = load_model_in_low_cpu_mem_usage(config, device_map, unused_kwargs, commit_hash, 
-                                                                      is_sharded, index_file)
-
-            else:
-                model = cls.from_config(config, **unused_kwargs)
-
-                state_dict = load_state_dict(model_file, variant=variant)
-                model._convert_deprecated_attn(state_dict)
-
-                model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
-                    model,
-                    state_dict,
-                    model_file,
-                    pretrained_model_name_or_path,
-                    ignore_mismatched_sizes=ignore_mismatched_sizes,
-                )
-
-                loading_info = {
-                    "missing_keys": missing_keys,
-                    "unexpected_keys": unexpected_keys,
-                    "mismatched_keys": mismatched_keys,
-                    "error_msgs": error_msgs,
-                }
-
-        if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
-            raise ValueError(
-                f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
-            )
-        elif torch_dtype is not None:
-            model = model.to(torch_dtype)
-    
-        model.register_to_config(_name_or_path=pretrained_model_name_or_path)
-
-        # Set model in evaluation mode to deactivate DropOut modules by default
-        model.eval()
-        if output_loading_info:
-            return model, loading_info
-
-        return model
-
-    @classmethod
-    def _load_pretrained_model(
-        cls,
-        model,
-        state_dict: OrderedDict,
-        resolved_archive_file,
-        pretrained_model_name_or_path: Union[str, os.PathLike],
-        ignore_mismatched_sizes: bool = False,
-    ):
-        # Retrieve missing & unexpected_keys
-        model_state_dict = model.state_dict()
-        loaded_keys = list(state_dict.keys())
-
-        expected_keys = list(model_state_dict.keys())
-
-        original_loaded_keys = loaded_keys
-
-        missing_keys = list(set(expected_keys) - set(loaded_keys))
-        unexpected_keys = list(set(loaded_keys) - set(expected_keys))
-
-        # Make sure we are able to load base models as well as derived models (with heads)
-        model_to_load = model
-
-        def _find_mismatched_keys(
-            state_dict,
-            model_state_dict,
-            loaded_keys,
-            ignore_mismatched_sizes,
-        ):
-            mismatched_keys = []
-            if ignore_mismatched_sizes:
-                for checkpoint_key in loaded_keys:
-                    model_key = checkpoint_key
-
-                    if (
-                        model_key in model_state_dict
-                        and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
-                    ):
-                        mismatched_keys.append(
-                            (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
-                        )
-                        del state_dict[checkpoint_key]
-            return mismatched_keys
-
-        if state_dict is not None:
-            # Whole checkpoint
-            mismatched_keys = _find_mismatched_keys(
-                state_dict,
-                model_state_dict,
-                original_loaded_keys,
-                ignore_mismatched_sizes,
-            )
-            error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
-
-        if len(error_msgs) > 0:
-            error_msg = "\n\t".join(error_msgs)
-            if "size mismatch" in error_msg:
-                error_msg += (
-                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
-                )
-            raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
-
-        if len(unexpected_keys) > 0:
-            logger.warning(
-                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
-                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
-                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task"
-                " or with another architecture (e.g. initializing a BertForSequenceClassification model from a"
-                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
-                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly"
-                " identical (initializing a BertForSequenceClassification model from a"
-                " BertForSequenceClassification model)."
-            )
-        else:
-            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
-        if len(missing_keys) > 0:
-            logger.warning(
-                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
-                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
-            )
-        elif len(mismatched_keys) == 0:
-            logger.info(
-                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the"
-                f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions"
-                " without further training."
-            )
-        if len(mismatched_keys) > 0:
-            mismatched_warning = "\n".join(
-                [
-                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
-                    for key, shape1, shape2 in mismatched_keys
-                ]
-            )
-            logger.warning(
-                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
-                f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be"
-                " able to use it for predictions and inference."
-            )
-
-        return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
-
-    @classmethod
-    def _get_signature_keys(cls, obj):
-        parameters = inspect.signature(obj.__init__).parameters
-        required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
-        optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
-        expected_modules = set(required_parameters.keys()) - {"self"}
-
-        return expected_modules, optional_parameters
-
-    # Adapted from `transformers` modeling_utils.py
-    def _get_no_split_modules(self, device_map: str):
-        """
-        Get the modules of the model that should not be spit when using device_map. We iterate through the modules to
-        get the underlying `_no_split_modules`.
-
-        Args:
-            device_map (`str`):
-                The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"]
-
-        Returns:
-            `List[str]`: List of modules that should not be split
-        """
-        _no_split_modules = set()
-        modules_to_check = [self]
-        while len(modules_to_check) > 0:
-            module = modules_to_check.pop(-1)
-            # if the module does not appear in _no_split_modules, we also check the children
-            if module.__class__.__name__ not in _no_split_modules:
-                if isinstance(module, ModelMixin):
-                    if module._no_split_modules is None:
-                        raise ValueError(
-                            f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement "
-                            f"support, the model class needs to implement the `_no_split_modules` attribute."
-                        )
-                    else:
-                        _no_split_modules = _no_split_modules | set(module._no_split_modules)
-                modules_to_check += list(module.children())
-        return list(_no_split_modules)
-
-    @property
-    def device(self) -> torch.device:
-        """
-        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
-        device).
-        """
-        return get_parameter_device(self)
-
-    @property
-    def dtype(self) -> torch.dtype:
-        """
-        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
-        """
-        return get_parameter_dtype(self)
-
-    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
-        """
-        Get number of (trainable or non-embedding) parameters in the module.
-
-        Args:
-            only_trainable (`bool`, *optional*, defaults to `False`):
-                Whether or not to return only the number of trainable parameters.
-            exclude_embeddings (`bool`, *optional*, defaults to `False`):
-                Whether or not to return only the number of non-embedding parameters.
-
-        Returns:
-            `int`: The number of parameters.
-
-        Example:
-
-        ```py
-        from diffusers import UNet2DConditionModel
-
-        model_id = "runwayml/stable-diffusion-v1-5"
-        unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
-        unet.num_parameters(only_trainable=True)
-        859520964
-        ```
-        """
-
-        if exclude_embeddings:
-            embedding_param_names = [
-                f"{name}.weight"
-                for name, module_type in self.named_modules()
-                if isinstance(module_type, torch.nn.Embedding)
-            ]
-            non_embedding_parameters = [
-                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
-            ]
-            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
-        else:
-            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
-
-    def _convert_deprecated_attn(self, state_dict: OrderedDict) -> None:
-        # original function name: _convert_deprecated_attention_blocks
-        deprecated_attention_block_paths = []
-
-        def recursive_find_attn_block(name, module):
-            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
-                deprecated_attention_block_paths.append(name)
-
-            for sub_name, sub_module in module.named_children():
-                sub_name = sub_name if name == "" else f"{name}.{sub_name}"
-                recursive_find_attn_block(sub_name, sub_module)
-
-        recursive_find_attn_block("", self)
-
-        # NOTE: we have to check if the deprecated parameters are in the state dict
-        # because it is possible we are loading from a state dict that was already
-        # converted
-
-        for path in deprecated_attention_block_paths:
-            # group_norm path stays the same
-
-            # query -> to_q
-            if f"{path}.query.weight" in state_dict:
-                state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")
-            if f"{path}.query.bias" in state_dict:
-                state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")
-
-            # key -> to_k
-            if f"{path}.key.weight" in state_dict:
-                state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")
-            if f"{path}.key.bias" in state_dict:
-                state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")
-
-            # value -> to_v
-            if f"{path}.value.weight" in state_dict:
-                state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")
-            if f"{path}.value.bias" in state_dict:
-                state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")
-
-            # proj_attn -> to_out.0
-            if f"{path}.proj_attn.weight" in state_dict:
-                state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")
-            if f"{path}.proj_attn.bias" in state_dict:
-                state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")
-
-
-class LegacyModelMixin(ModelMixin):
-    r"""
-    A subclass of `ModelMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
-    pipeline-specific classes (like `DiTTransformer2DModel`).
-    """
-
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        # To prevent dependency import problem.
-        from diffusers.models.model_loading_utils import _fetch_remapped_cls_from_config
-
-        # Create a copy of the kwargs so that we don't mess with the keyword arguments in the downstream calls.
-        kwargs_copy = kwargs.copy()
-
-        cache_dir = kwargs.pop("cache_dir", None)
-        force_download = kwargs.pop("force_download", False)
-        proxies = kwargs.pop("proxies", None)
-        local_files_only = kwargs.pop("local_files_only", None)
-        token = kwargs.pop("token", None)
-        revision = kwargs.pop("revision", None)
-        subfolder = kwargs.pop("subfolder", None)
-
-        # Load config if we don't provide a configuration
-        config_path = pretrained_model_name_or_path
-
-        user_agent = {
-            "diffusers": __version__,
-            "file_type": "model",
-            "framework": "pytorch",
-        }
-
-        # load config
-        config, _, _ = cls.load_config(
-            config_path,
-            cache_dir=cache_dir,
-            return_unused_kwargs=True,
-            return_commit_hash=True,
-            force_download=force_download,
-            proxies=proxies,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            subfolder=subfolder,
-            user_agent=user_agent,
-            **kwargs,
-        )
-        # resolve remapping
-        remapped_class = _fetch_remapped_cls_from_config(config, cls)
-
-        return remapped_class.from_pretrained(pretrained_model_name_or_path, **kwargs_copy)

omni_diffusion/__init__.py

@@ -1 +0,0 @@
-# omni_diffusion package - inference only, no eager imports

omni_diffusion/constants.py

@@ -1,105 +0,0 @@
-import logging
-
-logger = logging.getLogger(__name__)
-
-
-if False:
-    IMG_TAG_TOKEN = "<image>"
-    IMG_CONTEXT_TOKEN = "<IMG_CONTEXT>"
-    IMG_START_TOKEN = "<img>"
-    IMG_END_TOKEN = "</img>"
-
-    VID_TAG_TOKEN = "<video>"
-    VID_CONTEXT_TOKEN = "<VID_CONTEXT>"
-    VID_START_TOKEN = "<vid>"
-    VID_END_TOKEN = "</vid>"
-
-    PATCH_CONTEXT_TOKEN = "<PATCH_CONTEXT>"
-    PATCH_START_TOKEN = "<patch>"
-    PATCH_END_TOKEN = "</patch>"
-
-    AUD_TAG_TOKEN = "<audio>"
-    AUD_START_TOKEN = "<|begin_of_audio|>"
-    AUD_END_TOKEN = "<|end_of_audio|>"
-
-    QUAD_START_TOKEN = "<quad>"
-    QUAD_END_TOKEN = "</quad>"
-    REF_START_TOKEN = "<ref>"
-    REF_END_TOKEN = "</ref>"
-    BOX_START_TOKEN = "<box>"
-    BOX_END_TOKEN = "</box>"
-
-
-if True:
-
-    IMG_TAG_TOKEN = "<|image|>"
-    IMG_CONTEXT_TOKEN = "<|context_of_image|>"
-    IMG_START_TOKEN = "<|begin_of_image|>"
-    IMG_END_TOKEN = "<|end_of_image|>"
-
-    VID_TAG_TOKEN = "<|video|>"
-    VID_CONTEXT_TOKEN = "<|context_of_video|>"
-    VID_START_TOKEN = "<|begin_of_video|>"
-    VID_END_TOKEN = "<|end_of_video|>"
-
-    PATCH_CONTEXT_TOKEN = "<|context_of_patch|>"
-    PATCH_START_TOKEN = "<|begin_of_patch|>"
-    PATCH_END_TOKEN = "<|end_of_patch|>"
-
-    AUD_TAG_TOKEN = "<|audio|>"
-    AUD_CONTEXT_TOKEN = "<|context_of_audio|>"
-    AUD_START_TOKEN = "<|begin_of_audio|>"
-    AUD_END_TOKEN = "<|end_of_audio|>"
-
-    QUAD_START_TOKEN = "<|begin_of_quad|>"
-    QUAD_END_TOKEN = "<|end_of_quad|>"
-    REF_START_TOKEN = "<|begin_of_ref|>"
-    REF_END_TOKEN = "<|end_of_ref|>"
-    BOX_START_TOKEN = "<|begin_of_box|>"
-    BOX_END_TOKEN = "<|end_of_box|>"
-
-
-logger.info(f"{IMG_TAG_TOKEN=}")
-logger.info(f"{IMG_CONTEXT_TOKEN=}")
-logger.info(f"{IMG_START_TOKEN=}")
-logger.info(f"{IMG_END_TOKEN=}")
-
-logger.info(f"{VID_TAG_TOKEN=}")
-logger.info(f"{VID_CONTEXT_TOKEN=}")
-logger.info(f"{VID_START_TOKEN=}")
-logger.info(f"{VID_END_TOKEN=}")
-
-logger.info(f"{PATCH_CONTEXT_TOKEN=}")
-logger.info(f"{PATCH_START_TOKEN=}")
-logger.info(f"{PATCH_END_TOKEN=}")
-
-logger.info(f"{AUD_TAG_TOKEN=}")
-logger.info(f"{AUD_CONTEXT_TOKEN=}")
-logger.info(f"{AUD_START_TOKEN=}")
-logger.info(f"{AUD_END_TOKEN=}")
-
-# IMAGENET_MEAN = (0.485, 0.456, 0.406)
-# IMAGENET_STD = (0.229, 0.224, 0.225)
-
-# CLIP_MEAN = (0.4814546, 0.4578275, 0.40821073)
-# CLIP_STD = (0.2686295, 0.2613025, 0.2757711)
-
-# SIGLIP_MEAN = (0.5, 0.5, 0.5)
-# SIGLIP_STD = (0.5, 0.5, 0.5)
-
-
-IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406]
-IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225]
-IMAGENET_STANDARD_MEAN = [0.5, 0.5, 0.5]
-IMAGENET_STANDARD_STD = [0.5, 0.5, 0.5]
-OPENAI_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
-OPENAI_CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
-
-
-# Model Constants
-IGNORE_INDEX = -100
-IMAGE_TOKEN_INDEX = -200
-DEFAULT_IMAGE_TOKEN = IMG_CONTEXT_TOKEN
-DEFAULT_IMAGE_PATCH_TOKEN = PATCH_CONTEXT_TOKEN
-DEFAULT_IM_START_TOKEN = IMG_START_TOKEN
-DEFAULT_IM_END_TOKEN = IMG_END_TOKEN

omni_diffusion/data/__init__.py

@@ -1,2 +0,0 @@
-# data.__init__
-# Avoid importing build/dataset modules at package level for inference-only use

omni_diffusion/data/build.py

@@ -1,110 +0,0 @@
-import logging
-import os
-from dataclasses import dataclass
-
-from datasets import concatenate_datasets, load_dataset
-
-from .data_collator import DataCollatorForSupervisedDataset, collate_fn_deepspeed
-from .dataset_qwen2 import Qwen2Dataset
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-def build_supervised_dataset_deepspeed(
-    model_config,
-    model_args,
-    data_args,
-    training_args,
-    tokenizer,
-    create_position_ids=True,
-    create_loss_mask=False,
-    shift_token=False,
-):
-    """
-    Builds the supervised dataset and data collator specifically for DeepSpeed training.
-    
-    Args:
-        model_config: Configuration object for the model.
-        model_args: Arguments related to model architecture and parameters.
-        data_args: Arguments related to data processing and loading.
-        training_args: Arguments related to the training loop (e.g., output dir, seed).
-        tokenizer: The tokenizer instance to process text.
-        create_position_ids (bool): Whether to generate position IDs.
-        create_loss_mask (bool): Whether to generate a mask for loss calculation.
-        shift_token (bool): Whether to shift tokens for causal language modeling.
-
-    Returns:
-        dict: A dictionary containing 'train', 'validation' datasets and the 'data_collator'.
-    """
-    
-    logging.info("building dataset...")
-
-    cfg_path = data_args.dataset_name
-    max_padding_length = model_args.model_max_length
-    output_dir = training_args.output_dir
-
-    create_attention_mask = data_args.create_attention_mask
-    create_attention_mask_2d = data_args.create_attention_mask_2d
-
-    image_size = model_args.image_size
-    image_token_length = model_args.image_token_length
-
-    max_num_frame = model_args.max_num_frame
-    max_fps = model_args.max_fps
-
-    reset_position_ids = data_args.reset_position_ids
-    reset_attention_mask = data_args.reset_attention_mask
-    variable_length = data_args.variable_length
-
-    min_patch_grid = model_args.min_patch_grid
-    max_patch_grid = model_args.max_patch_grid
-    process_type = model_args.vision_process_type
-    normalize_type = model_args.vision_normalize_type
-
-    audio_tokenizer_path = model_args.audio_tokenizer_path
-    audio_tokenizer_type = model_args.audio_tokenizer_type
-    image_tokenizer_path = model_args.image_tokenizer_path
-
-    seed = training_args.seed
-    cross_dataset_joint = data_args.cross_dataset_joint
-    dataset_joint = data_args.dataset_joint
-
-    TrainDataset = Qwen2Dataset
-
-    train_dataset = TrainDataset(
-        cfg_path,
-        tokenizer,
-        image_size=image_size,
-        image_token_length=image_token_length,
-        max_padding_length=max_padding_length,
-        variable_length=variable_length,
-        output_dir=output_dir,
-        training_args=None,
-        shift_token=shift_token,
-        create_position_ids=create_position_ids,
-        create_attention_mask=create_attention_mask,
-        create_attention_mask_2d=create_attention_mask_2d,
-        create_loss_mask=create_loss_mask,
-        max_num_frame=max_num_frame,
-        max_fps=max_fps,
-        reset_position_ids=reset_position_ids,
-        reset_attention_mask=reset_attention_mask,
-        min_patch_grid=min_patch_grid,
-        max_patch_grid=max_patch_grid,
-        process_type=process_type,
-        normalize_type=normalize_type,
-        seed=seed,
-        cross_dataset_joint=cross_dataset_joint,
-        dataset_joint=dataset_joint,
-        audio_tokenizer_type=audio_tokenizer_type,
-        audio_tokenizer_path=audio_tokenizer_path,
-        image_tokenizer_path=image_tokenizer_path,
-        use_megatron=False,
-    )
-    eval_dataset = None
-
-    data_collator = collate_fn_deepspeed
-
-    return dict(train=train_dataset, validation=eval_dataset, data_collator=data_collator)
-

omni_diffusion/data/data_collator.py

@@ -1,122 +0,0 @@
-import itertools
-import json
-import logging
-import math
-import os
-import re
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Sequence
-
-import torch
-import transformers
-from torch.utils.data import default_collate
-from transformers.trainer_pt_utils import LabelSmoother
-
-IGNORE_TOKEN_ID = LabelSmoother.ignore_index
-
-
-@dataclass
-class DataCollatorForSupervisedDataset(object):
-    """
-    Data collator for supervised fine-tuning.
-    Handles padding of input IDs and labels, creation of attention masks,
-    and concatenation of multimodal inputs (images).
-    """
-    tokenizer: transformers.PreTrainedTokenizer
-
-    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
-        input_ids = [instance["input_ids"] for instance in instances]
-        labels = [instance["labels"] for instance in instances]
-
-        input_ids = torch.nn.utils.rnn.pad_sequence(
-            input_ids,
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id,
-        )
-        labels = torch.nn.utils.rnn.pad_sequence(
-            labels,
-            batch_first=True,
-            padding_value=IGNORE_TOKEN_ID,
-        )
-
-        input_ids = input_ids[:, : self.tokenizer.model_max_length]
-        labels = labels[:, : self.tokenizer.model_max_length]
-
-        batch = dict(
-            input_ids=input_ids,
-            labels=labels,
-            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
-        )
-
-        if "images" in instances[0]:
-            images = [instance["images"] for instance in instances]
-            batch["images"] = torch.cat(images, dim=0)
-
-        if "doclm_images" in instances[0]:
-            images = [instance["doclm_images"] for instance in instances]
-            batch["doclm_images"] = torch.cat(images, dim=0)
-
-        if "image_paths" in instances[0]:
-            image_paths = [instance["image_paths"] for instance in instances]
-            batch["image_paths"] = image_paths
-
-        if "pixel_values" in instances[0]:
-            pixel_values = torch.cat([instance["pixel_values"] for instance in instances])
-            batch["pixel_values"] = pixel_values
-
-        if "image_flags" in instances[0]:
-            image_flags = torch.cat([instance["image_flags"] for instance in instances])
-            batch["image_flags"] = image_flags
-
-        return batch
-
-def collate_fn_deepspeed(batch):
-    """
-    Custom collation function designed for DeepSpeed training.
-    Handles complex multimodal data (images/audio) and potentially packed sequences (cu_seq_lens).
-    """
-    tmp_batch = [{} for _ in range(len(batch))]
-    if "cu_seq_lens" in batch[0]:
-        cu_seq_lens = [x["cu_seq_lens"] for x in batch]
-        max_seq_len = [x["max_seq_len"] for x in batch]
-    else:
-        cu_seq_lens = None
-        max_seq_len = None
-
-    if "images" in batch[0].keys():
-        for new_x, x in zip(tmp_batch, batch):
-            new_x["images"] = x.pop("images")
-            new_x["image_indices"] = x.pop("image_indices")
-
-    if "audios" in batch[0].keys():
-        for new_x, x in zip(tmp_batch, batch):
-            new_x["audios"] = x.pop("audios")
-            new_x["audio_indices"] = x.pop("audio_indices")
-
-    new_batch = default_collate(batch)
-
-    if "images" in tmp_batch[0].keys():
-        
-        new_batch["images"] = torch.cat([x["images"] for x in tmp_batch], dim=0)
-
-        for sample_idx, sample in enumerate(tmp_batch):
-            sample["image_indices"][0, :, :] = sample_idx
-
-        new_batch["image_indices"] = torch.cat([x["image_indices"] for x in tmp_batch], dim=1)
-
-    if "audios" in tmp_batch[0].keys():
-        
-        new_batch["audios"] = list(itertools.chain.from_iterable([x["audios"] for x in tmp_batch]))
-        # print(f"{[x.size() for x in sample['audios']]}")
-
-        for sample_idx, sample in enumerate(tmp_batch):
-            for j in range(len(sample["audio_indices"])):
-                sample["audio_indices"][j][0, :, :] = sample_idx
-
-        new_batch["audio_indices"] = list(
-            itertools.chain.from_iterable([x["audio_indices"] for x in tmp_batch])
-        )
-        # print(f"{[x.size() for x in sample['audio_indices']]}")
-
-    
-    return new_batch

omni_diffusion/data/dataset_base.py

@@ -1,424 +0,0 @@
-import contextlib
-import json
-import logging
-import os
-import pdb
-import re
-import traceback
-import uuid
-
-import numpy as np
-import torch
-import yaml
-from PIL import Image
-
-from torchvision import transforms
-from torchvision.transforms import InterpolationMode
-
-from .processor.audio_processor import AudioProcessor
-from .processor.image_processor import ImageProcessor
-from .utils import draw_data
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-class BaseDataset(torch.utils.data.Dataset):
-    def __init__(
-        self,
-        cfg_path,
-        tokenizer,
-        image_size=448,
-        image_token_length=1024,
-        max_padding_length=32768,
-        variable_length=False,
-        output_dir="",
-        add_task_symbol=True,
-        training_args=None,
-        shift_token=False,
-        create_position_ids=True,
-        create_attention_mask=True,
-        create_attention_mask_2d=False,
-        create_loss_mask=False,
-        max_num_frame=8,
-        max_fps=1,
-        reset_position_ids=False,
-        reset_attention_mask=False,
-        min_patch_grid=1,
-        max_patch_grid=6,
-        process_type="anyres",
-        normalize_type="imagenet",
-        seed=42,
-        cross_dataset_joint=False,
-        dataset_joint=True,
-        audio_tokenizer_type=None,
-        audio_tokenizer_path=None,
-        image_tokenizer_path=None,
-        text_audio_interval_ratio=None,
-        use_megatron=True,
-        skip_samples=0,
-    ):
-        super(BaseDataset, self).__init__()
-
-        self.cfg_path = cfg_path
-        with open(self.cfg_path, "r", encoding="utf8") as cfg_file:
-            cfg_data = cfg_file.read()
-
-        self.cfg = yaml.load(cfg_data, Loader=yaml.CLoader)
-        logger.info(f"cfg {self.cfg}")
-
-        self.tokenizer = tokenizer
-        self.max_padding_length = max_padding_length
-        self.variable_length = variable_length
-        self.output_dir = output_dir
-        self.training_args = training_args
-        self.shift_token = shift_token
-        self.create_position_ids = create_position_ids
-        self.create_attention_mask = create_attention_mask
-        self.create_attention_mask_2d = create_attention_mask_2d
-        self.create_loss_mask = create_loss_mask
-        self.max_num_frame = max_num_frame
-        self.max_fps = max_fps
-        self.reset_position_ids = reset_position_ids
-        self.reset_attention_mask = reset_attention_mask
-
-        self.seed = seed
-        self.cross_dataset_joint = cross_dataset_joint
-        self.dataset_joint = dataset_joint
-
-        self.image_size = image_size
-        self.image_token_length = image_token_length
-
-        self.do_dataset_format = self.cfg.get("do_dataset_format", False)
-        self.do_dataset_cast = self.cfg.get("do_dataset_cast", False)
-        self.xlsx_sample_num = self.cfg.get("xlsx_sample_num", 5)
-
-        self.processor = {}
-        self.processor["image"] = ImageProcessor(
-            image_tokenizer_path,
-            process_type,
-            image_size=self.image_size,
-            normalize_type=normalize_type,
-            min_patch_grid=min_patch_grid,
-            max_patch_grid=max_patch_grid,
-        )
-
-        self.processor["audio"] = AudioProcessor(
-            audio_tokenizer_path=audio_tokenizer_path,
-            audio_tokenizer_type=audio_tokenizer_type,
-            text_audio_interval_ratio=text_audio_interval_ratio
-        )
-
-        self.total_samples = {}
-        self.processed_samples = {}
-        self.unjoint_samples = {}
-        self.joint_samples = {}
-        self.source2jsonpath = {}
-        self.source2prefixpath = {}
-        self.skip_samples = skip_samples
-
-        if use_megatron:
-            self.load_data()
-        else:
-            with main_process_first(local=False, desc="Loading data"):
-                self.load_data()
-
-
-    def load_data(self):
-        from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
-
-        raw_data = None
-
-        sampled_data = {}
-        source_idx = 0
-        for data_name, data_info in self.cfg["dataset"].items():
-            data_ratio = data_info.get("ratio", 1)
-            data_num = data_info.get("num", 999999999)
-
-            prefix_path = data_info.get("prefix_path", "")
-
-            if data_ratio == 0:
-                continue
-
-            if data_num == 0:
-                continue
-
-            for data_idx, data_path in enumerate(data_info["json_paths"]):
-                logger.info("-" * 100)
-
-                if not os.path.isfile(data_path) and not os.path.isdir(data_path):
-                    logger.warning(f"Data file no found {data_path}")
-                    continue
-
-                this_data = load_json(data_path, self.output_dir)
-                # this_data = load_data_one(data_path, self.outout_dir)
-                if this_data is None:
-                    logger.warning(f"Failed to load {data_path}")
-                    continue
-                # print(f"this_data {this_data}")
-
-                column_names = list(this_data.features)
-                if "id" in column_names:
-                    this_data = this_data.remove_columns("id")
-
-                # sources = [data_path] * len(this_data)
-                sources = [source_idx] * len(this_data)
-                # sources = [data_name] * len(this_data)
-                this_data = this_data.add_column("source", sources)
-
-                self.source2jsonpath[source_idx] = data_path
-                self.source2prefixpath[source_idx] = prefix_path
-                self.total_samples[source_idx] = len(this_data)
-                self.processed_samples[source_idx] = 0
-                self.unjoint_samples[source_idx] = 0
-                self.joint_samples[source_idx] = 0
-
-                source_idx += 1
-
-                if "images" not in column_names:
-                    # images = [[]] * len(this_data)
-                    images = [None] * len(this_data)
-                    this_data = this_data.add_column("images", images)
-
-                if "videos" not in column_names:
-                    # videos = [[]] * len(this_data)
-                    videos = [None] * len(this_data)
-                    this_data = this_data.add_column("videos", videos)
-
-                if "audios" not in column_names:
-                    # audios = [[]] * len(this_data)
-                    audios = [None] * len(this_data)
-                    this_data = this_data.add_column("audios", audios)
-
-                if False:
-                    column_names = list(this_data.features)
-                    this_data = this_data.map(
-                        format_function_general,
-                        batched=True,
-                        batch_size=2560,
-                        num_proc=1,
-                        # batch_size=1,
-                        # num_proc=1,
-                        remove_columns=column_names,
-                        keep_in_memory=False,
-                        desc="Running format on dataset",
-                    )
-
-                this_data = this_data.shuffle(seed=self.seed)
-                # this_data = this_data.flatten_indices()
-                this_data = this_data.shuffle(seed=self.seed)
-                # this_data = this_data.flatten_indices()
-
-                data_ratio = float(data_ratio)
-                total_num = len(this_data)
-                used_num = min(int(total_num * data_ratio), data_num)
-                logger.info(f"total_num {total_num}")
-                logger.info(f"data_ratio {data_ratio}")
-                logger.info(f"data_num {data_num}")
-                logger.info(f"used_num {used_num}")
-
-                indices = [x % total_num for x in range(used_num)]
-
-                this_data = this_data.select(indices)
-
-                if raw_data is None:
-                    raw_data = this_data
-                else:
-                    if self.do_dataset_cast:
-                        this_data = this_data.cast(raw_data.features)
-                    raw_data = concatenate_datasets([raw_data, this_data])
-
-                sampled_data[data_path] = {}
-                sampled_data[data_path]["data"] = this_data.select(
-                    range(min(self.xlsx_sample_num, used_num))
-                )
-                sampled_data[data_path]["total_num"] = total_num
-                sampled_data[data_path]["used_num"] = used_num
-
-                logger.info(f"this_data {this_data}")
-                logger.info(f"raw_data {raw_data}")
-                # logger.info(f"raw_data {raw_data[0]}")
-                # logger.info(f"raw_data {raw_data[-1]}")
-                logger.info(f"Successful load {data_path}")
-
-        raw_data = raw_data.shuffle(seed=self.seed)
-        # raw_data = raw_data.flatten_indices()
-        raw_data = raw_data.shuffle(seed=self.seed)
-        # raw_data = raw_data.flatten_indices()
-
-        self.raw_data = raw_data
-
-        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
-            output_xlsx = os.path.basename(self.cfg_path).replace("yaml", "xlsx")
-            output_xlsx = os.path.join(self.output_dir, output_xlsx)
-            logger.info(f"output_xlsx {output_xlsx}")
-            draw_data(
-                sampled_data,
-                output_xlsx,
-                tokenizer=self.tokenizer,
-                image_processor=self.processor["image"],
-            )
-
-        logger.info(f"raw_data {raw_data}")
-
-        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
-            logger.info(f"raw_data {raw_data[:10]}")
-            logger.info(f"raw_data {raw_data[-10:]}")
-
-    def __len__(self):
-        return len(self.raw_data)
-
-    def update_data_path(self, sample):
-        source = sample["source"]
-        prefix_path = self.source2prefixpath[source]
-
-        if has_image(sample):
-            sample["images"] = [os.path.join(prefix_path, image) for image in sample["images"]]
-
-        if has_video(sample):
-            sample["videos"] = [os.path.join(prefix_path, video) for video in sample["videos"]]
-
-        if has_audio(sample):
-            sample["audios"] = [os.path.join(prefix_path, audio) for audio in sample["audios"]]
-
-        return sample
-
-
-
-def format_function_general(examples):
-    messages = [x for x in examples["messages"]]
-
-    if "images" in examples:
-        images = [x for x in examples["images"]]
-    else:
-        images = [None for _ in messages]
-
-    if "videos" in examples:
-        videos = [x for x in examples["videos"]]
-    else:
-        videos = [None for _ in messages]
-
-    if "audios" in examples:
-        audios = [x for x in examples["audios"]]
-    else:
-        audios = [None for _ in messages]
-
-    return {
-        "messages": messages,
-        "images": images,
-        "videos": videos,
-        "audios": audios,
-    }
-
-
-def load_json_A(data_file):
-    from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
-
-    with open(data_file, "r") as f:
-        raw_data = json.load(f)
-    this_data = Dataset.from_list(raw_data)
-    return this_data
-
-
-def load_json_B(data_file):
-    from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
-
-    this_data = load_dataset("json", data_files=data_file, keep_in_memory=False)
-    return this_data["train"]
-
-
-def load_json_C(data_file):
-    from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
-
-    raw_data = []
-    with open(data_file, "r") as f:
-        for line in f.readlines():
-            d = json.loads(line)
-            # raw_data.append({"conversations": d["conversations"], "id": d["id"]})
-            if "conversations" in d:
-                raw_data.append({"conversations": d["conversations"]})
-            if "messages" in d:
-                raw_data.append({"messages": d["messages"]})
-    this_data = Dataset.from_list(raw_data)
-    return this_data
-
-
-def load_json(data_file, output_dir):
-    for func in [load_json_B, load_json_A, load_json_C]:
-        try:
-            this_data = func(data_file)
-            return this_data
-        except Exception as error:
-            with open(os.path.join(output_dir, "data_error.log"), "a") as f:
-                print("-" * 100, file=f)
-                print(f"{data_file=}", file=f)
-                # print(error, file=f)
-                print(traceback.format_exc(), file=f)
-            continue
-    return None
-
-
-def load_data_one(data_file, output_dir):
-    if data_file.endswith("json") or data_file.endswith("jsonl"):
-        return load_json(data_file, output_dir)
-    from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
-
-    this_data = load_dataset(data_file, keep_in_memory=False)
-    return this_data["train"]
-
-
-@contextlib.contextmanager
-def main_process_first(local=True, desc="work"):
-
-    if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
-        if local:
-            rank = int(os.environ["LOCAL_RANK"])
-        else:
-            rank = torch.distributed.get_rank()
-        is_main_process = rank == 0
-
-        try:
-            if not is_main_process:
-                torch.distributed.barrier()
-            yield
-        finally:
-            if is_main_process:
-                torch.distributed.barrier()
-    else:
-        yield
-
-def has_video(sample):
-    # video
-    if (
-        "videos" in sample
-        and isinstance(sample["videos"], list)
-        and None not in sample["videos"]
-        and len(sample["videos"])
-    ):
-        return True
-    return False
-
-
-def has_image(sample):
-    # image
-    if (
-        "images" in sample
-        and isinstance(sample["images"], list)
-        and None not in sample["images"]
-        and len(sample["images"])
-    ):
-        return True
-    return False
-
-
-def has_audio(sample):
-    # audio
-    if (
-        "audios" in sample
-        and isinstance(sample["audios"], list)
-        and None not in sample["audios"]
-        and len(sample["audios"])
-    ):
-        return True
-    return False
-

omni_diffusion/data/dataset_qwen2.py

@@ -1,981 +0,0 @@
-import json
-import logging
-import math
-import os
-import pdb
-import random
-import re
-import sys
-import time
-import traceback
-from collections import defaultdict
-from typing import Dict, List, Optional, Sequence
-from typing import List, Optional, Tuple, Union
-import copy
-import numpy as np
-import torch
-import transformers
-from transformers.trainer_pt_utils import LabelSmoother
-import soundfile as sf
-
-from .dataset_base import BaseDataset
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-IGNORE_TOKEN_ID = LabelSmoother.ignore_index
-
-def forward_process(
-    bsz: int,
-    seq_len: int,
-    device: torch.device,
-    labels: torch.Tensor,                
-    eps: float = 1e-3,
-    special_token_id: int = 151643,     
-    special_mask_ratio: float = 0.1      
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Generates a mask for the input sequences, applying different masking probabilities 
-    for normal tokens and special tokens within the valid (non-padded) regions.
-    """
-
-    # Initialize the boolean mask matrix and the overall mask ratio for each sample
-    b, l = bsz, seq_len
-    masked_indices = torch.zeros((b, l), device=device, dtype=torch.bool)
-    p_mask = torch.rand(b, device=device)
-    p_mask = (1 - eps) * p_mask + eps
-    p_mask = p_mask.unsqueeze(1)  # [b,1]
-
-    # Find the first and last valid positions (where label is not -100) for each sequence
-    first_idxs = []
-    last_idxs  = []
-    for i in range(b):
-        # -100 is the ignore_index
-        nonneg = (labels[i] != -100).nonzero(as_tuple=True)[0]
-        if nonneg.numel() == 0:
-            first_idxs.append(None)
-            last_idxs.append(None)
-        else:
-            first_idxs.append(int(nonneg[0]))
-            last_idxs.append(int(nonneg[-1]))
-
-    # Generate masks specifically for the valid interval of each sequence
-    for i in range(b):
-        start = first_idxs[i]
-        end   = last_idxs[i]
-
-        # Skip if the sequence has no valid tokens or invalid boundaries
-        if start is None or end is None or end < start:
-            continue
-
-        # Generate base thresholds for each position in the valid interval
-        valid_len = end - start + 1
-        t = torch.rand(valid_len, device=device)
-        mask_threshold = (1 - eps) * t + eps  # [valid_len]
-
-        # Generate random values to make masking decisions
-        rand_vals = torch.rand(valid_len, device=device)
-        # Normal token masking decision
-        normal_mask = rand_vals <= mask_threshold
-        # Special token masking threshold is lower
-        special_thresh = mask_threshold * special_mask_ratio
-        special_mask   = rand_vals <= special_thresh
-
-        labels_slice = labels[i, start : end + 1]
-        # Final mask: special tokens use special_mask, others use normal_mask
-        final_mask = torch.where(
-            labels_slice == special_token_id,
-            special_mask,
-            normal_mask
-        )
-
-        masked_indices[i, start : end + 1] = final_mask
-
-    total_masked   = int(masked_indices.sum().item())
-    special_masked = int((masked_indices & (labels == special_token_id)).sum().item())
-
-    return masked_indices, p_mask
-
-def update_labels(input_ids, labels, eos_id, max_n=20):
-    """
-    Finds the first occurrence of the EOS token in each sequence and updates 
-    up to `max_n` subsequent labels to the EOS token ID.
-    """
-    batch_size, seq_len = input_ids.shape
-    first_occurrence_indices = []
-
-    # Record the first occurrence position of eos_id in each batch sample
-    for idx in range(batch_size):
-        eos_positions = (input_ids[idx] == eos_id).nonzero(as_tuple=True)[0]
-        if len(eos_positions) > 0:
-            first_occurrence_indices.append(eos_positions[0].item())
-        else:
-            first_occurrence_indices.append(-1)  
-
-    # Select a random number of sequential positions (up to max_n) starting from first_idx to update
-    for i in range(batch_size):
-        first_idx = first_occurrence_indices[i]
-        if first_idx == -1:
-            continue 
-        max_possible = seq_len - first_idx
-
-        if max_possible <= 0:
-            continue
-        num_to_select = random.randint(1, min(max_n, max_possible))
-
-        selected_indices = torch.arange(first_idx, first_idx + num_to_select)
-
-        labels[i, selected_indices] = eos_id
-
-    return labels
-
-
-import torch
-import random
-
-def update_labels_and_inputs(input_ids, labels, eos_id, max_n=20):
-    input_ids = torch.tensor(input_ids).unsqueeze(0)
-    labels = torch.tensor(labels).unsqueeze(0)
-    batch_size, seq_len = input_ids.shape
-    input_ids = input_ids.clone()
-    labels = labels.clone()
-    new_input_ids = []
-    new_labels = []
-
-    for idx in range(batch_size):
-        eos_positions = (input_ids[idx] == eos_id).nonzero(as_tuple=True)[0]
-        if len(eos_positions) > 0:
-            first_idx = eos_positions[0].item()
-            cur_input_ids = input_ids[idx]
-            cur_labels = labels[idx]
-        else:
-            random_max_n = random.randint(1, max_n)
-            eos_ids = torch.full((random_max_n,), eos_id, device=input_ids.device, dtype=input_ids.dtype)
-            cur_input_ids = torch.cat([input_ids[idx], eos_ids])
-            pad_labels = torch.full((random_max_n,), eos_id, device=labels.device, dtype=labels.dtype)
-            cur_labels = torch.cat([labels[idx], pad_labels])
-            # first_idx = len(cur_input_ids) - random_max_n
-
-
-        new_input_ids.append(cur_input_ids)
-        new_labels.append(cur_labels)
-
-    max_len = max(len(x) for x in new_input_ids)
-    padded_input_ids = torch.stack([
-        torch.cat([x, torch.full((max_len - len(x),), eos_id, device=x.device, dtype=x.dtype)])
-        for x in new_input_ids
-    ])
-    padded_labels = torch.stack([
-        torch.cat([x, torch.full((max_len - len(x),), eos_id, device=x.device, dtype=x.dtype)])
-        for x in new_labels
-    ])
-
-    return padded_input_ids, padded_labels
-
-
-
-def pad_or_truncate_to_512(
-        input_ids, 
-        labels,
-        eos_id,
-        target_len: int = 512,
-):
-
-    input_ids = torch.as_tensor(input_ids).unsqueeze(0)   # (B, L)
-    labels    = torch.as_tensor(labels   ).unsqueeze(0)
-
-    batch_size = input_ids.size(0)
-    new_input_ids, new_labels = [], []
-
-    for i in range(batch_size):
-        cur_input = input_ids[i]
-        cur_label = labels[i]
-
-        cur_input = cur_input[:target_len]
-        cur_label = cur_label[:target_len]
-
-        pad_len = target_len - cur_input.size(0)
-        if pad_len > 0:
-            eos_pad   = torch.full((pad_len,), eos_id, device=cur_input.device, dtype=cur_input.dtype)
-            label_pad = torch.full((pad_len,), eos_id, device=cur_label.device, dtype=cur_label.dtype)
-            cur_input = torch.cat([cur_input, eos_pad])
-            cur_label = torch.cat([cur_label, label_pad])
-
-        new_input_ids.append(cur_input)
-        new_labels.append(cur_label)
-
-    return torch.stack(new_input_ids), torch.stack(new_labels)
-
-
-
-class Qwen2Dataset(BaseDataset):
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
-        super().__init__(
-            *args,
-            **kwargs,
-        )
-
-        self.default_system_message = "You are a helpful AI assistant."
-        self.default_system_message = None
-
-        self.ret = defaultdict(dict)
-        self.is_cat = True
-
-        if self.cross_dataset_joint:
-            for i in range(2):
-                self.maybe_init_ret(f"default_{i}")
-
-    def maybe_init_ret(self, source, force=False):
-        """Initializes the return dictionary buffer for a specific source if it doesn't exist."""
-        if source not in self.ret or force:
-            self.ret[source] = {}
-
-            self.ret[source]["tokens"] = []
-            self.ret[source]["labels"] = []
-            self.ret[source]["cu_seq_lens"] = [0]
-            self.ret[source]["max_seq_len"] = self.max_padding_length
-
-            if self.create_position_ids:
-                self.ret[source]["position_ids"] = []
-
-            if self.create_attention_mask:
-                self.ret[source]["attention_mask"] = []
-
-            if self.create_attention_mask_2d:
-                self.ret[source]["attention_mask_2d"] = torch.tril(
-                    torch.ones(
-                        (1, self.max_padding_length, self.max_padding_length), dtype=torch.bool
-                    )
-                )
-        return len(self.ret[source]["tokens"]) == 0
-
-    def get_max_min_ret_length(self):
-        """Finds the buffer with the maximum and minimum current lengths."""
-        max_ret_lengh = 0
-        min_ret_lengh = self.max_padding_length + 1
-
-        max_ret_key = None
-        min_ret_key = None
-
-        for k, v in self.ret.items():
-            cur_length = len(v["tokens"])
-
-            if cur_length > max_ret_lengh:
-                max_ret_lengh = cur_length
-                max_ret_key = k
-
-            if cur_length < min_ret_lengh:
-                min_ret_lengh = cur_length
-                min_ret_key = k
-
-        return max_ret_lengh, max_ret_key, min_ret_lengh, min_ret_key
-
-    def add_ret(self, ret, source):
-        """
-        Appends a processed sample (tokens, images, audios) to the buffer (`self.ret`).
-        Updates indices for images/audios to account for the offset in the concatenated sequence.
-        """
-        cur_length = len(ret["input_ids"])
-        cur_image_length = len(ret["images"])
-        cur_audio_length = len(ret["audios"])
-
-        all_length = len(self.ret[source]["tokens"])
-
-        if "images" in self.ret[source]:
-            all_image_length = len(self.ret[source]["images"])
-        else:
-            all_image_length = 0
-
-        if cur_image_length > 0:
-            if all_image_length > 0:
-                self.ret[source]["images"] = torch.cat(
-                    [self.ret[source]["images"], ret["images"]], dim=0
-                )
-                ret["image_indices"][1, :, :] += all_length
-                self.ret[source]["image_indices"] = torch.cat(
-                    [self.ret[source]["image_indices"], ret["image_indices"]], dim=1
-                )
-            else:
-                self.ret[source]["images"] = ret["images"]
-                self.ret[source]["image_indices"] = ret["image_indices"]
-
-        if "audios" in self.ret[source]:
-            all_audio_length = len(self.ret[source]["audios"])
-        else:
-            all_audio_length = 0
-
-        if cur_audio_length > 0:
-            if all_audio_length > 0:
-                self.ret[source]["audios"].extend(ret["audios"])
-                for audio_indice in ret["audio_indices"]:
-                    audio_indice[1, :, :] += all_length
-                self.ret[source]["audio_indices"].extend(ret["audio_indices"])
-            else:
-                self.ret[source]["audios"] = ret["audios"]
-                self.ret[source]["audio_indices"] = ret["audio_indices"]
-
-            # print(self.ret[source]["audios"])
-
-        if self.create_attention_mask:
-            self.ret[source]["attention_mask"] += ret["attention_mask"]
-
-        if self.create_attention_mask_2d:
-            self.ret[source]["attention_mask_2d"][:, all_length:, :all_length] = 0
-
-        if self.create_position_ids:
-            self.ret[source]["position_ids"] += list(range(cur_length))
-
-        self.ret[source]["tokens"] += ret["input_ids"]
-        self.ret[source]["labels"] += ret["labels"]
-        self.ret[source]["cu_seq_lens"] += [all_length + cur_length]
-
-    def process_ret(self, to_ret):
-        """
-        Finalizes the buffer for return. Handles padding, truncation, tensor conversion,
-        and attention mask generation.
-        """
-        if "tokens" in to_ret and len(to_ret["tokens"]) > 0:
-            pass
-        else:
-            return to_ret
-
-        if self.create_position_ids:
-            if self.reset_position_ids:
-                pass
-            else:
-                to_ret["position_ids"] = list(range(len(to_ret["tokens"])))
-
-        if self.create_attention_mask_2d:
-            if self.reset_attention_mask:
-                pass
-            else:
-                to_ret["attention_mask_2d"] = torch.tril(
-                    torch.ones(
-                        (1, self.max_padding_length, self.max_padding_length), dtype=torch.bool
-                    )
-                )
-
-        if self.shift_token:
-            to_ret["tokens"] = to_ret["tokens"][:-1]
-            to_ret["labels"] = to_ret["labels"][1:]
-            to_ret["cu_seq_lens"][-1] -= 1
-            if self.create_position_ids:
-                to_ret["position_ids"] = to_ret["position_ids"][:-1]
-            if self.create_attention_mask:
-                to_ret["attention_mask"] = to_ret["attention_mask"][:-1]
-
-            if self.create_attention_mask_2d:
-                to_ret["attention_mask_2d"][:, :, -1] = 0
-                to_ret["attention_mask_2d"][:, -1, :] = 0
-
-        assert len(to_ret["tokens"]) == len(
-            to_ret["labels"]
-        ), f"{len(to_ret['tokens'])} {len(to_ret['labels'])}"
-
-        if not self.variable_length and self.max_padding_length > len(to_ret["tokens"]):
-            to_ret["tokens"] += [self.tokenizer.pad_token_id] * (
-                self.max_padding_length - len(to_ret["tokens"])
-            )
-            to_ret["labels"] += [IGNORE_TOKEN_ID] * (
-                self.max_padding_length - len(to_ret["labels"])
-            )
-            to_ret["cu_seq_lens"][-1] = self.max_padding_length
-            if self.create_position_ids:
-                # to_ret["position_ids"] += to_ret["position_ids"][-1:] * (
-                #     self.max_padding_length - len(to_ret["position_ids"])
-                # )
-                to_ret["position_ids"] += list(
-                    range(to_ret["position_ids"][-1] + 1, self.max_padding_length)
-                )
-            if self.create_attention_mask:
-                to_ret["attention_mask"] += [0] * (
-                    self.max_padding_length - len(to_ret["attention_mask"])
-                )
-
-        to_ret["tokens"] = to_ret["tokens"][: self.max_padding_length]
-        to_ret["labels"] = to_ret["labels"][: self.max_padding_length]
-        to_ret["cu_seq_lens"][-1] = self.max_padding_length
-
-        if self.create_position_ids:
-            to_ret["position_ids"] = to_ret["position_ids"][: self.max_padding_length]
-        if self.create_attention_mask:
-            to_ret["attention_mask"] = to_ret["attention_mask"][: self.max_padding_length]
-
-        to_ret["tokens"] = torch.tensor(to_ret["tokens"], dtype=torch.int64)
-        to_ret["labels"] = torch.tensor(to_ret["labels"], dtype=torch.int64)
-        to_ret["cu_seq_lens"] = torch.tensor(to_ret["cu_seq_lens"], dtype=torch.int64)
-        if self.create_position_ids:
-            to_ret["position_ids"] = torch.tensor(to_ret["position_ids"], dtype=torch.int64)
-        if self.create_attention_mask:
-            to_ret["attention_mask"] = torch.tensor(to_ret["attention_mask"], dtype=torch.int64)
-
-        if self.create_attention_mask_2d:
-            attention_mask_2d = to_ret.pop("attention_mask_2d")
-            attention_mask_2d = attention_mask_2d.masked_fill(
-                (to_ret["attention_mask"] < 0.5).view(1, 1, self.max_padding_length), value=0
-            )
-            attention_mask_2d = attention_mask_2d < 0.5
-
-            to_ret["attention_mask"] = attention_mask_2d
-
-        if self.create_loss_mask:
-            loss_mask = torch.where(to_ret["labels"] == IGNORE_TOKEN_ID, 0, 1)
-            to_ret["loss_mask"] = loss_mask.to(torch.float32)
-
-        if not self.reset_position_ids and not self.reset_attention_mask:
-            to_ret.pop("cu_seq_lens")
-        else:
-            max_seq_len = max(to_ret["cu_seq_lens"][1:] - to_ret["cu_seq_lens"][:-1])
-            to_ret["max_seq_len"] = max_seq_len
-        to_ret["input_ids"] = to_ret["tokens"]
-
-        return to_ret
-
-    def is_skip(self):
-
-        processed_samples = sum(self.processed_samples.values())
-        if processed_samples < self.skip_samples:
-            if processed_samples % 1e3 == 0:
-                print(
-                    f"processed_samples {processed_samples} skip_samples {self.skip_samples}"
-                )
-            return True
-
-    def show_statistic(self):
-        log_interval = 2000
-        if self.max_padding_length >= 2**17:
-            log_interval = 1000
-        if self.max_padding_length >= 2**20:
-            log_interval = 200 
-
-        processed_samples = sum(self.processed_samples.values())
-        unjoint_samples = sum(self.unjoint_samples.values())
-        joint_samples = sum(self.joint_samples.values())
-        if unjoint_samples % log_interval == 1:
-            pass
-        else:
-            return
-
-        with open(os.path.join(self.output_dir, "data_statistics.log"), "a") as f:
-            print("-" * 100, file=f)
-            print(
-                f"processed_samples {processed_samples}" +
-                f" unjoint_samples {unjoint_samples}" +
-                f" joint_samples {joint_samples}" +
-                f" {[len(v['tokens']) for _, v in self.ret.items()]}", file=f,
-            )
-
-            print("source processed_samples  unjoint_samples    joint_samples data_path", file=f)
-            for source, data_path in self.source2jsonpath.items():
-                print(f"{source: >6}  {self.processed_samples[source]: >16} \
-                      {self.unjoint_samples[source]: >16} {self.joint_samples[source]: >16} {data_path}", file=f)
-
-    def __getitem__(self, index):
-        """
-        Main data loading function.
-        1. Fetches raw data.
-        2. Preprocesses it (tokenization, multimodal handling).
-        3. Packs multiple samples into a single buffer (`self.ret`) until `max_padding_length` is reached.
-        4. Returns the packed batch when full.
-        """
-        self.processor["audio"].load_model()
-        self.processor["image"].load_model()
-
-        while True:
-            try:
-                sample = self.raw_data[index]
-                sample = copy.deepcopy(sample)
-                sample = self.update_data_path(sample)
-                source = sample["source"]
-
-                self.processed_samples[source] += 1
-                if self.is_skip():
-                    return {}
-
-                if self.cross_dataset_joint:
-                    is_empty = False
-                    (
-                        max_ret_lengh,
-                        max_ret_key,
-                        min_ret_lengh,
-                        min_ret_key,
-                    ) = self.get_max_min_ret_length()
-                else:
-                    is_empty = self.maybe_init_ret(source)
-
-                    max_ret_lengh = min_ret_lengh = len(self.ret[source]["tokens"])
-                    max_ret_key = min_ret_key = source
-
-                is_begin = is_empty or self.reset_position_ids or self.reset_attention_mask
-
-                #logger.info("preprocess begin " + str(self.processor["audio"].audio_tokenizer.device))
-
-                ret = preprocess(
-                    sample,
-                    self.tokenizer,
-                    self.image_token_length,
-                    default_system_message=self.default_system_message,
-                    processor=self.processor,
-                    is_begin=is_begin,
-                    max_num_frame=self.max_num_frame,
-                    max_fps=self.max_fps,
-                )
-
-                #logger.info("preprocess end " + str(self.processor["audio"].audio_tokenizer.device))
-
-                if ret is None:
-                    #logger.info("ret is None " + str(self.processor["audio"].audio_tokenizer.device))
-                    return {}
-
-                cur_length = len(ret["input_ids"])
-
-                if cur_length > self.max_padding_length:
-                    return {}
-
-                self.unjoint_samples[source] += 1
-
-                if not self.dataset_joint:
-                    import pdb; pdb.set_trace()
-                    to_ret = self.ret.pop(max_ret_key)
-
-                    self.maybe_init_ret(max_ret_key, force=True)
-                    self.add_ret(ret, max_ret_key)
-
-                elif min_ret_lengh + cur_length > self.max_padding_length:
-                    #logger.info("data too long " + str(self.processor["audio"].audio_tokenizer.device))
-                    to_ret = self.ret.pop(max_ret_key)
-                    self.joint_samples[source] += 1
-
-                    self.maybe_init_ret(max_ret_key, force=True)
-                    self.add_ret(ret, max_ret_key)
-
-                else:
-                    to_ret = {}
-                    self.add_ret(ret, min_ret_key)
-
-                to_ret = self.process_ret(to_ret)
-
-                self.show_statistic()
-                #logger.info("output ret " + str(len(ret)) + " " + str(self.processor["audio"].audio_tokenizer.device))
-                return to_ret
-
-            except Exception as error:
-                try:
-                    with open(os.path.join(self.output_dir, "data_error.log"), "a") as f:
-                        #print("-" * 100)
-                        #print(traceback.format_exc())
-                        #print(self.raw_data[index])
-                        print("-" * 100, file=f)
-                        print(traceback.format_exc(), file=f)
-                        print(self.raw_data[index], file=f)
-                except Exception as error:
-                    print(error)
-                return {}
-
-
-def preprocess(
-    sample,
-    tokenizer: transformers.PreTrainedTokenizer,
-    image_token_length: int,
-    default_system_message: str = "You are a helpful assistant.",
-    processor=None,
-    is_begin: bool = True,
-    max_num_frame: int = 8,
-    max_fps: int = 1,
-) -> Dict:
-
-    from ..constants import (
-        IMG_START_TOKEN,
-        IMG_END_TOKEN,
-        IMG_CONTEXT_TOKEN,
-        VID_START_TOKEN,
-        VID_END_TOKEN,
-        VID_CONTEXT_TOKEN,
-        PATCH_START_TOKEN,
-        PATCH_END_TOKEN,
-        PATCH_CONTEXT_TOKEN,
-        AUD_START_TOKEN,
-        AUD_END_TOKEN,
-        IMG_TAG_TOKEN,
-        VID_TAG_TOKEN,
-        AUD_TAG_TOKEN,
-        AUD_CONTEXT_TOKEN,
-    )
-
-    human_roles = ["user", "human"]
-    gpt_roles = ["assistant", "gpt"]
-    system_roles = ["system"]
-
-    # Ensure special tokens map to exactly one ID
-    IMG_CONTEXT_ID = tokenizer(IMG_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    IMG_START_ID = tokenizer(IMG_START_TOKEN, add_special_tokens=False).input_ids
-    IMG_END_ID = tokenizer(IMG_END_TOKEN, add_special_tokens=False).input_ids
-
-    VID_CONTEXT_ID = tokenizer(VID_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    VID_START_ID = tokenizer(VID_START_TOKEN, add_special_tokens=False).input_ids
-    VID_END_ID = tokenizer(VID_END_TOKEN, add_special_tokens=False).input_ids
-
-    PATCH_CONTEXT_ID = tokenizer(PATCH_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    PATCH_START_ID = tokenizer(PATCH_START_TOKEN, add_special_tokens=False).input_ids
-    PATCH_END_ID = tokenizer(PATCH_END_TOKEN, add_special_tokens=False).input_ids
-
-    AUD_CONTEXT_ID = tokenizer(AUD_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    AUD_START_ID = tokenizer(AUD_START_TOKEN, add_special_tokens=False).input_ids
-    AUD_END_ID = tokenizer(AUD_END_TOKEN, add_special_tokens=False).input_ids
-
-    IMG_TAG_ID = tokenizer(IMG_TAG_TOKEN, add_special_tokens=False).input_ids
-    VID_TAG_ID = tokenizer(VID_TAG_TOKEN, add_special_tokens=False).input_ids
-    AUD_TAG_ID = tokenizer(AUD_TAG_TOKEN, add_special_tokens=False).input_ids
-
-    assert len(IMG_CONTEXT_ID) == 1
-    assert len(IMG_START_ID) == 1
-    assert len(IMG_END_ID) == 1
-
-    assert len(VID_CONTEXT_ID) == 1
-    assert len(VID_START_ID) == 1
-    assert len(VID_END_ID) == 1
-
-    assert len(PATCH_CONTEXT_ID) == 1
-    assert len(PATCH_START_ID) == 1
-    assert len(PATCH_END_ID) == 1
-
-    IMG_CONTEXT_ID = IMG_CONTEXT_ID[0]
-    IMG_START_ID = IMG_START_ID[0]
-    IMG_END_ID = IMG_END_ID[0]
-
-    VID_CONTEXT_ID = VID_CONTEXT_ID[0]
-    VID_START_ID = VID_START_ID[0]
-    VID_END_ID = VID_END_ID[0]
-
-    PATCH_CONTEXT_ID = PATCH_CONTEXT_ID[0]
-    PATCH_START_ID = PATCH_START_ID[0]
-    PATCH_END_ID = PATCH_END_ID[0]
-
-    AUD_CONTEXT_ID = AUD_CONTEXT_ID[0]
-    AUD_START_ID = AUD_START_ID[0]
-    AUD_END_ID = AUD_END_ID[0]
-
-    IMG_TAG_ID = IMG_TAG_ID[0]
-    VID_TAG_ID = VID_TAG_ID[0]
-    AUD_TAG_ID = AUD_TAG_ID[0]
-
-    BOS_ID = tokenizer.bos_token_id
-    EOS_ID = tokenizer.eos_token_id
-
-    # ChatML format special tokens
-    IM_START = "<|im_start|>"
-    IM_END = "<|im_end|>"
-    USER = "user"
-    ASSISTANT = "assistant"
-    SYSTEM = "system"
-
-    nl_tokens = tokenizer("\n", add_special_tokens=False).input_ids
-    IM_START_IDS = tokenizer(IM_START, add_special_tokens=False).input_ids
-    IM_END_IDS = tokenizer(IM_END, add_special_tokens=False).input_ids
-    USER_IDS = tokenizer(USER, add_special_tokens=False).input_ids
-    ASSISTANT_IDS = tokenizer(ASSISTANT, add_special_tokens=False).input_ids
-    SYSTEM_IDS = tokenizer(SYSTEM, add_special_tokens=False).input_ids
-
-    input_ids, targets = [], []
-    images = []
-    image_indices = []
-    audios = []
-    audio_indices = []
-
-    messages = []
-    if "conversations" in sample:
-        messages = sample["conversations"]
-    if len(messages) == 0 and "messages" in sample:
-        messages = sample["messages"]
-
-    # ----------------------------------------------------------------
-    # 1. System Prompt Processing
-    # ----------------------------------------------------------------
-    has_system = False
-    if is_begin:
-        if messages[0]["role"] == "system":
-            has_system = True
-        else:
-            has_system = False
-
-        if (
-            not has_system
-            and default_system_message is not None
-            and len(default_system_message) > 0
-        ):
-            messages = [{"role": "system", "content": default_system_message}] + messages
-            has_system = True
-
-    # ----------------------------------------------------------------
-    # 2. Audio Processing
-    # ----------------------------------------------------------------
-    if has_audio(sample) and processor["audio"].is_discrete:
-        unused_audio_idxs = list(range(len(sample["audios"])))
-        
-        audio_tokens_list = [
-            processor["audio"].process_audios(x, is_discrete=True) for x in sample["audios"]
-        ]
-        audio_tokens_list = ["".join(f"<|audio_{i}|>" for i in x) for x in audio_tokens_list]
-
-        audio_idx = 0
-        for j, sentence in enumerate(messages):
-            content = sentence["content"]
-            role = sentence["role"]
-            # whether apply discrete tokenize to this role
-            if processor["audio"].apply_to_role(role, is_discrete=True):
-                while AUD_TAG_TOKEN in content:
-                    content = content.replace(
-                        AUD_TAG_TOKEN,
-                        f"{AUD_START_TOKEN}{audio_tokens_list[audio_idx]}{AUD_END_TOKEN}",
-                        1,
-                    )
-                    # <|begin_of_audio|> <|audio_0|> <|audio_1|> ... <|audio_n|> <|end_of_audio|>
-                    unused_audio_idxs.remove(audio_idx)
-                    audio_idx += 1
-            else:
-                audio_idx += content.count(AUD_TAG_TOKEN)
-
-            sentence["content"] = content
-
-    # ----------------------------------------------------------------
-    # 3. Image Processing
-    # ----------------------------------------------------------------
-    if has_image(sample):
-        # for visual question answering & captioning
-        image_tokens_512_list = [
-            processor["image"].process_images_with_subpatch(x, 512) for x in sample["images"]
-        ]
-        image_tokens_512_list = [
-            processor["image"].get_image_token(x) for x in image_tokens_512_list
-        ]
-        image_tokens_512_list = [x[0].tolist() for x in image_tokens_512_list]
-        image_tokens_512_list = ["".join(f"<|image_{i}|>" for i in x) for x in image_tokens_512_list]
-
-        # for image generation
-        image_tokens_256_list = [
-            processor["image"].process_images_with_subpatch(x, 256) for x in sample["images"]
-        ]
-        image_tokens_256_list = [
-            processor["image"].get_image_token(x) for x in image_tokens_256_list
-        ]
-        image_tokens_256_list = [x[0].tolist() for x in image_tokens_256_list]
-        image_tokens_256_list = ["".join(f"<|image_{i}|>" for i in x) for x in image_tokens_256_list]
-
-        image_idx = 0
-        for j, sentence in enumerate(messages):
-            content = sentence["content"]
-            role = sentence["role"]
-            # for image, always apply discrete tokenize to this role
-            # if processor["image"].apply_to_role(role) or True:
-            if role == "user":
-                image_resolution = 512
-            else:
-                image_resolution = 256
-
-            while IMG_TAG_TOKEN in content:
-                if image_resolution == 256:
-                    content = content.replace(
-                        IMG_TAG_TOKEN,
-                        f"{IMG_START_TOKEN}{image_tokens_256_list[image_idx]}{IMG_END_TOKEN}",
-                        1,
-                    )
-                    # <|begin_of_image|> <|image_0|> <|image_1|> ... <|image_n|> <|end_of_image|>
-                else:
-                    content = content.replace(
-                        IMG_TAG_TOKEN,
-                        f"{IMG_START_TOKEN}{image_tokens_512_list[image_idx]}{IMG_END_TOKEN}",
-                        1,
-                    )
-                    # <|begin_of_image|> <|image_0|> <|image_1|> ... <|image_n|> <|end_of_image|>
-                image_idx += 1
-            else:
-                image_idx += content.count(IMG_TAG_TOKEN)
-
-            sentence["content"] = content
-
-    # ----------------------------------------------------------------
-    # 4. Text Processing
-    # ----------------------------------------------------------------
-    for j, sentence in enumerate(messages):
-        role = sentence["role"]
-        content = sentence["content"]
-
-        if role in human_roles:
-            _input_id = (
-                IM_START_IDS
-                + USER_IDS
-                + nl_tokens
-                + tokenizer(content, add_special_tokens=False).input_ids
-                + IM_END_IDS
-                + nl_tokens
-            )
-            _target = [IGNORE_TOKEN_ID] * len(_input_id)
-
-        elif role in gpt_roles:
-            content_input_id = tokenizer(content, add_special_tokens=False).input_ids
-
-            _input_id = (
-                IM_START_IDS + ASSISTANT_IDS + nl_tokens + content_input_id + IM_END_IDS + nl_tokens
-            )
-            _target = (
-                [IGNORE_TOKEN_ID] * len(IM_START_IDS)
-                + [IGNORE_TOKEN_ID] * len(ASSISTANT_IDS)
-                + [IGNORE_TOKEN_ID] * len(nl_tokens)
-                + content_input_id
-                + IM_END_IDS
-                + nl_tokens
-            )
-
-            dream_pad_token = tokenizer.encode("<|endoftext|>")
-            _input_id_dream = (
-                IM_START_IDS + ASSISTANT_IDS + nl_tokens + content_input_id + IM_END_IDS + dream_pad_token
-            )
-            _target_dream = (
-                [IGNORE_TOKEN_ID] * len(IM_START_IDS)
-                + [IGNORE_TOKEN_ID] * len(ASSISTANT_IDS)
-                + [IGNORE_TOKEN_ID] * len(nl_tokens)
-                + content_input_id
-                + IM_END_IDS
-                + dream_pad_token
-            )
-
-        elif role in system_roles:
-            _input_id = (
-                IM_START_IDS
-                + SYSTEM_IDS
-                + nl_tokens
-                + tokenizer(content, add_special_tokens=False).input_ids
-                + IM_END_IDS
-                + nl_tokens
-            )
-            _target = [IGNORE_TOKEN_ID] * len(_input_id)
-
-        else:
-            raise NotImplementedError
-
-        input_ids += _input_id
-        targets += _target
-
-
-    # ----------------------------------------------------------------
-    # 5. Contiguous Audio Processing
-    # ----------------------------------------------------------------
-    if has_audio(sample) and processor["audio"].is_contiguous:
-        aud_positions = [i for i, x in enumerate(input_ids) if x == AUD_TAG_ID]
-        assert len(aud_positions) == len(unused_audio_idxs), sample
-
-        new_input_ids = []
-        new_targets = []
-        st = 0
-
-        for aud_idx, aud_pos in enumerate(aud_positions):
-            aud_idx = unused_audio_idxs[aud_idx]
-            audio = processor["audio"].process_audios(sample["audios"][aud_idx], is_contiguous=True)
-            audios.append(audio)
-            audio_token_length = audio.size(0) + 4
-            # audio_token_length = audio.size(0)
-
-            new_input_ids += input_ids[st:aud_pos]
-            new_targets += targets[st:aud_pos]
-
-            new_input_ids += [AUD_START_ID]
-            new_targets += [IGNORE_TOKEN_ID]
-
-            audio_indice_b = torch.zeros(
-                1, audio_token_length, dtype=torch.int64
-            )  # This will change in collate_fn
-            audio_indice_s = (
-                torch.arange(len(new_input_ids), len(new_input_ids) + audio_token_length)
-                .unsqueeze(0)
-                .repeat(1, 1)
-            )
-            audio_indice_b_s = torch.stack(
-                [audio_indice_b, audio_indice_s], dim=0
-            )  # 2, num_image, image_length
-            audio_indices.append(audio_indice_b_s)
-
-            new_input_ids += [AUD_CONTEXT_ID] * audio_token_length
-            new_targets += [IGNORE_TOKEN_ID] * audio_token_length
-
-            new_input_ids += [AUD_END_ID]
-            new_targets += [IGNORE_TOKEN_ID]
-
-            st = aud_pos + 1
-
-        new_input_ids += input_ids[st:]
-        new_targets += targets[st:]
-
-        input_ids = new_input_ids
-        targets = new_targets
-
-    if len(images) > 0:
-        images = torch.cat(images, dim=0)
-
-    if len(image_indices) > 0:
-        image_indices = torch.cat(image_indices, dim=1)
-
-    # ----------------------------------------------------------------
-    # 6. Final Masking and Label Updating 
-    # ----------------------------------------------------------------
-    origin_input = input_ids
-    labels = targets
-    eos_id = 151643               
-    mask_id = 151666    
-    input_ids, labels = update_labels_and_inputs(input_ids,labels,eos_id,16)
-
-    labels_mask           = ~(labels == -100) 
-    bsz, seq_len          = labels_mask.shape
-    masked_indices, p_mask = forward_process(
-        bsz, seq_len, input_ids.device, labels,special_mask_ratio=0.6, special_token_id=eos_id
-    )
-    final_masked_indices = masked_indices & labels_mask
-    final_masked_indices_inv = (~masked_indices) & labels_mask
-
-    mask_id_tensor = torch.full_like(input_ids, mask_id)
-    input_ids = torch.where(final_masked_indices, mask_id_tensor, input_ids)
-
-    new_labels = labels.clone()
-    new_labels[final_masked_indices_inv] = -100
-
-    input_ids = input_ids.squeeze(0).cpu().tolist()
-    new_labels = new_labels.squeeze(0).cpu().tolist()
-    attention_mask = [1] * len(input_ids)
-    assert len(new_labels) == len(input_ids)
-    
-    return dict(
-        input_ids=input_ids,
-        labels=new_labels,
-        attention_mask=attention_mask,
-        images=images,
-        image_indices=image_indices,
-        audios=audios,
-        audio_indices=audio_indices,
-    )
-
-
-def has_image(sample):
-    # image
-    if (
-        "images" in sample
-        and isinstance(sample["images"], list)
-        and None not in sample["images"]
-        and len(sample["images"])
-    ):
-        return True
-    return False
-
-
-def has_audio(sample):
-    # audio
-    if (
-        "audios" in sample
-        and isinstance(sample["audios"], list)
-        and None not in sample["audios"]
-        and len(sample["audios"])
-    ):
-        return True
-    return False

omni_diffusion/data/processor/audio_processor.py

@@ -1,142 +0,0 @@
-import json
-import math
-import os
-
-import numpy as np
-import torch
-
-import natsort
-from omni_diffusion.tokenizer import get_audio_tokenizer
-
-
-class AudioProcessor:
-    """
-    Handles the initialization of audio tokenizers and processing of audio files.
-    """
-    def __init__(
-        self,
-        audio_tokenizer_path=None,
-        audio_tokenizer_type=None,
-        text_audio_interval_ratio=None,
-    ):
-
-        self.audio_tokenizer = get_audio_tokenizer(
-            audio_tokenizer_path,
-            audio_tokenizer_type,
-        )
-
-        self.audio_tokenizer_type = audio_tokenizer_type
-
-        self.text_audio_interval_ratio = text_audio_interval_ratio
-
-        # self.load_model()
-
-    def load_model(self):
-        """Loads the underlying tokenizer model weights."""
-        if self.audio_tokenizer is not None:
-            self.audio_tokenizer.load_model()
-
-    def process_audios(self, audio_path, is_discrete=False, is_contiguous=False, **kwargs):
-
-        assert not (is_discrete and is_contiguous)
-        assert is_discrete or is_contiguous
-
-        if is_discrete:
-            audio_tokenizer_type = self.audio_tokenizer_type.split("_")[-1]
-            cache_path = os.path.splitext(audio_path)[0] + f"_{audio_tokenizer_type}.json"
-            try:
-                if os.path.isfile(cache_path):
-                    with open(cache_path, "r") as f:
-                        audio_data = json.load(f)
-                    return audio_data
-            except Exception as e:
-                pass
-
-        audio_data = self.audio_tokenizer.encode(
-            audio_path, is_discrete=is_discrete, is_contiguous=is_contiguous, **kwargs
-        )
-        # print(f"{len(audio_data)=}")
-
-        if is_discrete:
-            try:
-                if isinstance(audio_data, list):
-                    with open(cache_path, "w") as f:
-                        json.dump(audio_data, f)
-            except Exception as e:
-                pass
-
-        return audio_data
-
-    @property
-    def is_discrete(self):
-        return self.audio_tokenizer.is_discrete
-
-    @property
-    def is_contiguous(self):
-        return self.audio_tokenizer.is_contiguous
-
-    def apply_to_role(self, role, **kwargs):
-        return self.audio_tokenizer.apply_to_role(role, **kwargs)
-
-
-def add_audio_input_contiguous(input_ids, audio_paths, tokenizer, audio_tokenizer):
-    """
-    Injects contiguous audio embeddings into a text sequence.
-    Replaces specific audio tag tokens in the input_ids with a sequence of 
-    [AUD_START] + [AUD_CONTEXT]... + [AUD_END].
-    """
-    from ...constants import (
-        AUD_START_TOKEN,
-        AUD_END_TOKEN,
-        AUD_TAG_TOKEN,
-        AUD_CONTEXT_TOKEN,
-    )
-
-    AUD_CONTEXT_ID = tokenizer(AUD_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    AUD_TAG_ID = tokenizer(AUD_TAG_TOKEN, add_special_tokens=False).input_ids
-    AUD_START_ID = tokenizer(AUD_START_TOKEN, add_special_tokens=False).input_ids
-    AUD_END_ID = tokenizer(AUD_END_TOKEN, add_special_tokens=False).input_ids
-
-    AUD_CONTEXT_ID = AUD_CONTEXT_ID[0]
-    AUD_TAG_ID = AUD_TAG_ID[0]
-    AUD_START_ID = AUD_START_ID[0]
-    AUD_END_ID = AUD_END_ID[0]
-
-    aud_positions = [i for i, x in enumerate(input_ids) if x == AUD_TAG_ID]
-
-    audios = []
-    audio_indices = []
-    new_input_ids = []
-    st = 0
-    for aud_idx, aud_pos in enumerate(aud_positions):
-        audio = audio_tokenizer.encode(audio_paths[aud_idx], is_contiguous=True)
-        audios.append(audio)
-        audio_token_length = audio.size(0) + 4
-
-        new_input_ids += input_ids[st:aud_pos]
-
-        new_input_ids += [AUD_START_ID]
-
-        audio_indice_b = torch.zeros(
-            1, audio_token_length, dtype=torch.int64
-        )  # This will change in collate_fn
-        audio_indice_s = (
-            torch.arange(len(new_input_ids), len(new_input_ids) + audio_token_length)
-            .unsqueeze(0)
-            .repeat(1, 1)
-        )
-        audio_indice_b_s = torch.stack(
-            [audio_indice_b, audio_indice_s], dim=0
-        )  # 2, num_image, image_length
-        audio_indices.append(audio_indice_b_s)
-
-        new_input_ids += [AUD_CONTEXT_ID] * audio_token_length
-
-        new_input_ids += [AUD_END_ID]
-
-        st = aud_pos + 1
-
-    new_input_ids += input_ids[st:]
-    inputs_ids = new_input_ids
-
-    return inputs_ids, audios, audio_indices

omni_diffusion/data/processor/image_processor.py

@@ -1,358 +0,0 @@
-import math
-import os
-
-import numpy as np
-import torch
-from PIL import Image
-
-import decord
-import natsort
-from omni_diffusion.constants import (
-    IMAGENET_DEFAULT_MEAN,
-    IMAGENET_DEFAULT_STD,
-    IMAGENET_STANDARD_MEAN,
-    IMAGENET_STANDARD_STD,
-    OPENAI_CLIP_MEAN,
-    OPENAI_CLIP_STD,
-)
-from torchvision import transforms
-from omni_diffusion.tokenizer_magvitv2 import MagVITV2Tokenizer
-
-import logging
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-def image_transform(image, resolution=256):
-    image = transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BICUBIC)(image)
-    image = transforms.CenterCrop((resolution, resolution))(image)
-    image = transforms.ToTensor()(image)
-    image = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)(image)
-    return image
-
-
-class ImageProcessor:
-    def __init__(
-        self,
-        model_path,
-        process_type,
-        image_size=256,
-        normalize_type="imagenet",
-        min_patch_grid=1,
-        max_patch_grid=6,
-    ):
-        self.process_type = process_type
-        self.image_size = image_size
-
-        if normalize_type == "imagenet":
-            MEAN, STD = IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-        elif normalize_type == "clip":
-            MEAN, STD = OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
-        elif normalize_type == "siglip":
-            MEAN, STD = IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD
-        else:
-            raise NotImplementedError(normalize_type)
-        self.mean = MEAN
-        self.std = STD
-
-        self.patch_size = image_size
-        self.min_patch_grid = min_patch_grid
-        self.max_patch_grid = max_patch_grid
-
-        if self.process_type == "anyres":
-            self.grid_pinpoints = [
-                (i, j)
-                for i in range(min_patch_grid, max_patch_grid + 1)
-                for j in range(min_patch_grid, max_patch_grid + 1)
-            ]
-            self.possible_resolutions = [
-                [dim * self.patch_size for dim in pair] for pair in self.grid_pinpoints
-            ]
-            print(f"grid_pinpoints {self.grid_pinpoints}")
-            print(f"possible_resolutions {self.possible_resolutions}")
-
-        if self.process_type == "dynamic":
-            max_num = self.max_patch_grid
-            min_num = self.min_patch_grid
-            # calculate the existing image aspect ratio
-            target_ratios = set(
-                (i, j)
-                for n in range(min_num, max_num + 1)
-                for i in range(1, n + 1)
-                for j in range(1, n + 1)
-                if i * j <= max_num and i * j >= min_num
-            )
-            self.target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
-            self.possible_resolutions = [
-                [dim * self.patch_size for dim in pair] for pair in self.target_ratios
-            ]
-            print(f"target_ratios {self.target_ratios}")
-            print(f"possible_resolutions {self.possible_resolutions}")
-
-        # self.image_tokenizer = MAGVITv2()
-        self.image_tokenizer = MagVITV2Tokenizer(model_path=model_path)
-
-    def load_model(self):
-        if self.image_tokenizer is not None:
-            self.image_tokenizer.load_model()
-        
-    def process_images(self, img_or_path_list, image_resolution):
-
-        if isinstance(img_or_path_list[0], str):
-            images = [Image.open(x).convert("RGB") for x in img_or_path_list]
-        elif isinstance(img_or_path_list[0], Image.Image):
-            images = [x.convert("RGB") for x in img_or_path_list]
-        else:
-            images = img_or_path_list
-
-        # image_tensor = torch.ones([len(images), 3, self.image_size, self.image_size])
-        image_tensor = torch.ones([len(images), 3, image_resolution, image_resolution])
-
-        for i, image in enumerate(images):
-            image = image_transform(image, resolution=image_resolution)
-            
-            image_tensor[i] = image
-
-        return image_tensor
-
-    def process_images_with_subpatch(self, img_or_path, image_resolution):
-        return self.process_images([img_or_path], image_resolution)
-    
-    def get_image_token(self, image):
-        return self.image_tokenizer.encode(image)
-    
-
-def select_best_resolution(original_size, possible_resolutions):
-    """
-    Selects the best resolution from a list of possible resolutions based on the original size.
-
-    Args:
-        original_size (tuple): The original size of the image in the format (width, height).
-        possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), 
-            (width2, height2), ...].
-
-    Returns:
-        tuple: The best fit resolution in the format (width, height).
-    """
-    original_width, original_height = original_size
-    best_fit = None
-    max_effective_resolution = 0
-    min_wasted_resolution = float("inf")
-
-    for width, height in possible_resolutions:
-        # Calculate the downscaled size to keep the aspect ratio
-        scale = min(width / original_width, height / original_height)
-        downscaled_width, downscaled_height = int(original_width * scale), int(
-            original_height * scale
-        )
-
-        # Calculate effective and wasted resolutions
-        effective_resolution = min(
-            downscaled_width * downscaled_height, original_width * original_height
-        )
-        wasted_resolution = (width * height) - effective_resolution
-
-        if effective_resolution > max_effective_resolution or (
-            effective_resolution == max_effective_resolution
-            and wasted_resolution < min_wasted_resolution
-        ):
-            max_effective_resolution = effective_resolution
-            min_wasted_resolution = wasted_resolution
-            best_fit = (width, height)
-
-    return best_fit
-
-
-def resize_and_pad_image(image, target_resolution):
-    """
-    Resize and pad an image to a target resolution while maintaining aspect ratio.
-
-    Args:
-        image (PIL.Image.Image): The input image.
-        target_resolution (tuple): The target resolution (width, height) of the image.
-
-    Returns:
-        PIL.Image.Image: The resized and padded image.
-    """
-    original_width, original_height = image.size
-    target_width, target_height = target_resolution
-
-    # Determine which dimension (width or height) to fill
-    scale_w = target_width / original_width
-    scale_h = target_height / original_height
-
-    if scale_w < scale_h:
-        # Width will be filled completely
-        new_width = target_width
-        new_height = min(math.ceil(original_height * scale_w), target_height)
-    else:
-        # Height will be filled completely
-        new_height = target_height
-        new_width = min(math.ceil(original_width * scale_h), target_width)
-
-    # Resize the image
-    resized_image = image.resize((new_width, new_height))
-
-    # Create a new image with the target size and paste the resized image onto it
-    new_image = Image.new("RGB", (target_width, target_height), (0, 0, 0))
-    paste_x = (target_width - new_width) // 2
-    paste_y = (target_height - new_height) // 2
-    new_image.paste(resized_image, (paste_x, paste_y))
-
-    return new_image
-
-
-def add_image_input_contiguous(input_ids, image_paths, tokenizer):
-
-    image_processor = ImageProcessor(
-        process_type="dynamic",
-        image_size=448,
-        normalize_type="imagenet",
-        min_patch_grid=1,
-        max_patch_grid=12,
-    )
-
-    image_token_length = 256
-    max_num_frame = 4096
-    max_fps = 1
-
-    from ...constants import (
-        IMG_START_TOKEN,
-        IMG_END_TOKEN,
-        IMG_TAG_TOKEN,
-        IMG_CONTEXT_TOKEN,
-        VID_START_TOKEN,
-        VID_END_TOKEN,
-        VID_TAG_TOKEN,
-        VID_CONTEXT_TOKEN,
-        PATCH_START_TOKEN,
-        PATCH_END_TOKEN,
-        PATCH_CONTEXT_TOKEN,
-    )
-
-    IMG_CONTEXT_ID = tokenizer(IMG_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    IMG_START_ID = tokenizer(IMG_START_TOKEN, add_special_tokens=False).input_ids
-    IMG_END_ID = tokenizer(IMG_END_TOKEN, add_special_tokens=False).input_ids
-
-    VID_CONTEXT_ID = tokenizer(VID_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    VID_START_ID = tokenizer(VID_START_TOKEN, add_special_tokens=False).input_ids
-    VID_END_ID = tokenizer(VID_END_TOKEN, add_special_tokens=False).input_ids
-
-    PATCH_CONTEXT_ID = tokenizer(PATCH_CONTEXT_TOKEN, add_special_tokens=False).input_ids
-    PATCH_START_ID = tokenizer(PATCH_START_TOKEN, add_special_tokens=False).input_ids
-    PATCH_END_ID = tokenizer(PATCH_END_TOKEN, add_special_tokens=False).input_ids
-
-    IMG_TAG_ID = tokenizer(IMG_TAG_TOKEN, add_special_tokens=False).input_ids
-    VID_TAG_ID = tokenizer(VID_TAG_TOKEN, add_special_tokens=False).input_ids
-
-    assert len(IMG_CONTEXT_ID) == 1
-    assert len(IMG_START_ID) == 1
-    assert len(IMG_END_ID) == 1
-
-    assert len(VID_CONTEXT_ID) == 1
-    assert len(VID_START_ID) == 1
-    assert len(VID_END_ID) == 1
-
-    assert len(PATCH_CONTEXT_ID) == 1
-    assert len(PATCH_START_ID) == 1
-    assert len(PATCH_END_ID) == 1
-
-    IMG_CONTEXT_ID = IMG_CONTEXT_ID[0]
-    IMG_START_ID = IMG_START_ID[0]
-    IMG_END_ID = IMG_END_ID[0]
-
-    VID_CONTEXT_ID = VID_CONTEXT_ID[0]
-    VID_START_ID = VID_START_ID[0]
-    VID_END_ID = VID_END_ID[0]
-
-    PATCH_CONTEXT_ID = PATCH_CONTEXT_ID[0]
-    PATCH_START_ID = PATCH_START_ID[0]
-    PATCH_END_ID = PATCH_END_ID[0]
-
-    IMG_TAG_ID = IMG_TAG_ID[0]
-    VID_TAG_ID = VID_TAG_ID[0]
-
-    nl_tokens = tokenizer("\n", add_special_tokens=False).input_ids
-
-    img_positions = [i for i, x in enumerate(input_ids) if x == IMG_TAG_ID]
-
-    images = []
-    image_indices = []
-    new_input_ids = []
-    st = 0
-    for img_idx, img_pos in enumerate(img_positions):
-        image_patches, (
-            best_width,
-            best_height,
-        ) = image_processor.process_images_with_subpatch(image_paths[img_idx])
-        images.append(image_patches)
-        print(f"add_image_input_contiguous best_width {best_width} best_height {best_height}")
-
-        new_input_ids += input_ids[st:img_pos]
-
-        new_input_ids += [IMG_START_ID]
-
-        image_indice_b = torch.zeros(
-            1, image_token_length, dtype=torch.int64
-        )  # This will change in collate_fn
-        image_indice_s = (
-            torch.arange(len(new_input_ids), len(new_input_ids) + image_token_length)
-            .unsqueeze(0)
-            .repeat(1, 1)
-        )
-        image_indice_b_s = torch.stack(
-            [image_indice_b, image_indice_s], dim=0
-        )  # 2, num_image, image_length
-        image_indices.append(image_indice_b_s)
-
-        new_input_ids += [IMG_CONTEXT_ID] * image_token_length
-
-        new_input_ids += [IMG_END_ID]
-
-        if len(image_patches) > 1:
-            for i in range(0, best_height, image_processor.patch_size):
-                new_input_ids += nl_tokens
-
-                for j in range(0, best_width, image_processor.patch_size):
-                    new_input_ids += [PATCH_START_ID]
-
-                    image_indice_b = torch.zeros(
-                        1, image_token_length, dtype=torch.int64
-                    )  # This will change in collate_fn
-                    image_indice_s = (
-                        torch.arange(
-                            len(new_input_ids), len(new_input_ids) + image_token_length
-                        )
-                        .unsqueeze(0)
-                        .repeat(1, 1)
-                    )
-                    image_indice_b_s = torch.stack(
-                        [image_indice_b, image_indice_s], dim=0
-                    )  # 2, num_image, image_length
-                    image_indices.append(image_indice_b_s)
-
-                    new_input_ids += [PATCH_CONTEXT_ID] * image_token_length
-
-                    new_input_ids += [PATCH_END_ID]
-                    # print(f"get_external_dict i {i} j {j} new_input_ids {len(new_input_ids)}")
-
-        st = img_pos + 1
-
-    new_input_ids += input_ids[st:]
-
-    inputs_ids = new_input_ids
-
-    images = torch.cat(images, dim=0)
-    image_indices = torch.cat(image_indices, dim=1)
-
-    image_indices = image_indices.contiguous().to(torch.cuda.current_device())
-    if True:
-        images = (
-            torch.tensor(images, dtype=torch.bfloat16).contiguous().to(torch.cuda.current_device())
-        )
-
-    else:
-        images = (
-            torch.tensor(images, dtype=torch.float16).contiguous().to(torch.cuda.current_device())
-        )
-
-    return inputs_ids, images, image_indices

omni_diffusion/data/utils.py

@@ -1,212 +0,0 @@
-import argparse
-import copy
-import hashlib
-import io
-import itertools
-import json
-import logging
-import os
-import re
-from typing import Dict, Sequence, Tuple
-
-import numpy as np
-import xlsxwriter
-from datasets import concatenate_datasets, load_dataset
-from PIL import Image
-from tabulate import tabulate
-from termcolor import colored
-
-
-def buffer_image(image: Image, format: str = "JPEG"):
-    # Store image in buffer, so we don't have to write it to disk.
-    buffer = io.BytesIO()
-    image.save(buffer, format=format)
-    return buffer, image
-
-
-def resize(img_or_path: str, size: Tuple[int, int], format="JPEG"):
-    """
-    Opens an image (from a file path or existing object), creates a thumbnail 
-    preserving aspect ratio using LANCZOS filter, and returns the image buffer.
-    """
-    if isinstance(img_or_path, str):
-        image = Image.open(img_or_path)
-    else:
-        image = img_or_path
-    # image = image.resize(size)
-    image.thumbnail(size, Image.LANCZOS)
-    image = image.convert("RGB")
-
-    return buffer_image(image, format)
-
-
-def calculate_scale(file_path, bound_size):
-    # check the image size without loading it into memory
-    im = Image.open(file_path)
-    original_width, original_height = im.size
-
-    # calculate the resize factor, keeping original aspect and staying within boundary
-    bound_width, bound_height = bound_size
-    ratios = (float(bound_width) / original_width, float(bound_height) / original_height)
-    return min(ratios)
-
-
-def draw_data(all_datasets, output_path, tokenizer=None, image_processor=None):
-    """
-    Visualizes dataset contents by exporting them to an Excel file (.xlsx).
-    It creates a summary sheet and individual sheets for each dataset, 
-    embedding text, images, and video frames into the cells.
-    """
-    if hasattr(tokenizer, "image_start_tag"):
-        image_start_tag = tokenizer.image_start_tag
-        image_end_tag = tokenizer.image_end_tag
-    else:
-        from ..constants import IMG_START_TOKEN, IMG_END_TOKEN
-
-        image_start_tag = IMG_START_TOKEN
-        image_end_tag = IMG_END_TOKEN
-
-    # Initialize Excel Workbook
-    workbook = xlsxwriter.Workbook(output_path)
-    cell_format = workbook.add_format({"text_wrap": True, "font_size": 12})
-
-    worksheet = workbook.add_worksheet("ALL")
-    worksheet.set_column(1, 2, 20, cell_format)
-    worksheet.set_column(3, 3, 240, cell_format)
-    worksheet.write(0, 1, "total_num")
-    worksheet.write(0, 2, "used_num")
-    worksheet.write(0, 3, "name")
-    row = 1
-
-    # Write statistics for each dataset
-    for this_name, this_dataset in all_datasets.items():
-        total_num = this_dataset["total_num"]
-        used_num = this_dataset["used_num"]
-        worksheet.write(row, 1, total_num)
-        worksheet.write(row, 2, used_num)
-        worksheet.write(row, 3, this_name)
-        row += 1
-    worksheet.write(
-        row, 2, sum([this_dataset["used_num"] for this_dataset in all_datasets.values()])
-    )
-
-    all_base_name = ["all"]
-    for this_name, this_dataset in all_datasets.items():
-        base_name = os.path.basename(this_name)
-        base_name = os.path.splitext(base_name)[0]
-
-        base_name = base_name[:24]
-        all_base_name.append(base_name)
-
-        if all_base_name.count(base_name) > 1:
-            base_name = base_name + "_" + str(all_base_name.count(base_name))
-
-        worksheet = workbook.add_worksheet(base_name)
-        worksheet.set_column(1, 2, 120, cell_format)
-        worksheet.write(0, 1, "user")
-        worksheet.write(0, 2, "assistant")
-        row = 1
-
-        data = this_dataset["data"]
-        for this_data in data:
-            # print(this_data)
-            if isinstance(this_data, Dict):
-                # print(this_data.keys())
-                messages = this_data["messages"]
-                if "images" in this_data:
-                    images = this_data["images"]
-                if "videos" in this_data:
-                    videos = this_data["videos"]
-            else:
-                messages = this_data
-
-            image_count = 0
-            video_count = 0
-            for message in messages:
-                content = message["content"]
-                role = message["role"]
-                if role == "user" or role == "human":
-                    col = 1
-                else:
-                    col = 2
-                worksheet.write(row, col, content)
-                row += 1
-
-                bos_pos = [m.start() for m in re.finditer(image_start_tag, content)]
-                eos_pos = [m.start() for m in re.finditer(image_end_tag, content)]
-                # print(bos_pos, eos_pos)
-                for a, b in zip(bos_pos, eos_pos):
-                    # print(content[a+len(image_start_tag:b])
-                    img_path = content[a + len(image_start_tag) : b]
-                    # print(img_path)
-                    worksheet.set_row(row, 200)
-
-                    try:
-                        image_buffer, image = resize(img_path, (512, 512), format="JPEG")
-                    except:
-                        continue
-
-                    scale = min(256 / image.width, 256 / image.height)
-                    data = {"x_scale": scale, "y_scale": scale, "object_position": 1}
-
-                    worksheet.insert_image(row, col, img_path, {"image_data": image_buffer, **data})
-
-                    row += 1
-
-                for _ in range(content.count("<image>") + content.count("<|image|>")):
-                    if images is None:
-                        continue
-                    img_path = images[image_count]
-                    # print(img_path)
-                    worksheet.set_row(row, 200)
-
-                    try:
-                        image_buffer, image = resize(img_path, (512, 512), format="JPEG")
-                    except:
-                        continue
-
-                    scale = min(256 / image.width, 256 / image.height)
-                    data = {"x_scale": scale, "y_scale": scale, "object_position": 1}
-
-                    worksheet.insert_image(row, col, img_path, {"image_data": image_buffer, **data})
-
-                    row += 1
-                    image_count += 1
-
-                for _ in range(content.count("<video>") + content.count("<|video|>")):
-                    if videos is None:
-                        continue
-                    vid_path = videos[video_count]
-                    try:
-                        _, video_frames = image_processor.process_video(vid_path, max_num_frame=4)
-                        # print(vid_path)
-                    except:
-                        continue
-
-                    for video_frame in video_frames:
-                        worksheet.set_row(row, 200)
-                        try:
-                            image_buffer, image = resize(video_frame, (512, 512), format="JPEG")
-                        except:
-                            continue
-
-                        scale = min(256 / image.width, 256 / image.height)
-                        data = {"x_scale": scale, "y_scale": scale, "object_position": 1}
-
-                        if isinstance(video_frame, str):
-                            video_path = video_frame
-                        else:
-                            video_file = hashlib.md5(video_frame.tobytes()).hexdigest() + ".png"
-                            video_path = os.path.join("/tmp/", video_file)
-                            video_frame.save(video_path)
-
-                        worksheet.insert_image(
-                            row, col, video_path, {"image_data": image_buffer, **data}
-                        )
-
-                        row += 1
-                    video_count += 1
-
-            row += 8
-
-    workbook.close()

omni_diffusion/models/__init__.py

@@ -1,2 +0,0 @@
-# models.__init__
-# Avoid eager imports to prevent loading training-only dependencies

omni_diffusion/models/dream/__init__.py

@@ -1,12 +0,0 @@
-# models.dream.__init__
-
-from .modeling_dream import DreamModel
-from .configuration_dream import DreamConfig  
-from .tokenization_dream import DreamTokenizer
-from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForCausalLM
-
-AutoConfig.register("Dream", DreamConfig)
-AutoModelForCausalLM.register(DreamConfig, DreamModel)
-
-DreamConfig.register_for_auto_class()
-DreamModel.register_for_auto_class("AutoModelForCausalLM")

omni_diffusion/models/dream/config.yaml

@@ -1,98 +0,0 @@
-encoder: SenseVoiceEncoderSmall
-encoder_conf:
-    output_size: 512
-    attention_heads: 4
-    linear_units: 2048
-    num_blocks: 50
-    tp_blocks: 20
-    dropout_rate: 0.1
-    positional_dropout_rate: 0.1
-    attention_dropout_rate: 0.1
-    input_layer: pe
-    pos_enc_class: SinusoidalPositionEncoder
-    normalize_before: true
-    kernel_size: 11
-    sanm_shfit: 0
-    selfattention_layer_type: sanm
-
-
-model: SenseVoiceSmall
-model_conf:
-    length_normalized_loss: true
-    sos: 1
-    eos: 2
-    ignore_id: -1
-
-tokenizer: SentencepiecesTokenizer
-tokenizer_conf:
-  bpemodel: null
-  unk_symbol: <unk>
-  split_with_space: true
-
-frontend: WavFrontend
-frontend_conf:
-    fs: 16000
-    window: hamming
-    n_mels: 80
-    frame_length: 25
-    frame_shift: 10
-    lfr_m: 7
-    lfr_n: 6
-    cmvn_file: null
-
-
-dataset: SenseVoiceCTCDataset
-dataset_conf:
-  index_ds: IndexDSJsonl
-  batch_sampler: EspnetStyleBatchSampler
-  data_split_num: 32
-  batch_type: token
-  batch_size: 14000
-  max_token_length: 2000
-  min_token_length: 60
-  max_source_length: 2000
-  min_source_length: 60
-  max_target_length: 200
-  min_target_length: 0
-  shuffle: true
-  num_workers: 4
-  sos: ${model_conf.sos}
-  eos: ${model_conf.eos}
-  IndexDSJsonl: IndexDSJsonl
-  retry: 20
-
-train_conf:
-  accum_grad: 1
-  grad_clip: 5
-  max_epoch: 20
-  keep_nbest_models: 10
-  avg_nbest_model: 10
-  log_interval: 100
-  resume: true
-  validate_interval: 10000
-  save_checkpoint_interval: 10000
-
-optim: adamw
-optim_conf:
-  lr: 0.00002
-scheduler: warmuplr
-scheduler_conf:
-  warmup_steps: 25000
-
-specaug: SpecAugLFR
-specaug_conf:
-    apply_time_warp: false
-    time_warp_window: 5
-    time_warp_mode: bicubic
-    apply_freq_mask: true
-    freq_mask_width_range:
-    - 0
-    - 30
-    lfr_rate: 6
-    num_freq_mask: 1
-    apply_time_mask: true
-    time_mask_width_range:
-    - 0
-    - 12
-    num_time_mask: 1
-

omni_diffusion/models/dream/config_dream_resume.json

@@ -1,37 +0,0 @@
-{
-    "_name_or_path": "Instruct-7B",
-    "architectures": [
-      "DreamModel"
-    ],
-    "attention_dropout": 0.0,
-    "auto_map": {
-      "AutoConfig": "configuration_dream.DreamConfig",
-      "AutoModel": "modeling_dream.DreamModel"
-    },
-    "bos_token_id": 151643,
-    "eos_token_id": 151643,
-    "hidden_act": "silu",
-    "hidden_size": 3584,
-    "initializer_range": 0.02,
-    "intermediate_size": 18944,
-    "mask_token_id": 151666,
-    "max_position_embeddings": 131072,
-    "max_window_layers": 28,
-    "model_type": "Dream",
-    "num_attention_heads": 28,
-    "num_hidden_layers": 28,
-    "num_key_value_heads": 4,
-    "pad_token_id": 151643,
-    "rms_norm_eps": 1e-06,
-    "rope_scaling": null,
-    "rope_theta": 1000000.0,
-    "sliding_window": null,
-    "tie_word_embeddings": false,
-    "torch_dtype": "bfloat16",
-    "transformers_version": "4.46.2",
-    "use_cache": true,
-    "use_mrope": false,
-    "use_sliding_window": false,
-    "vocab_size": 176264,
-    "chunk_size": -1
-  }

omni_diffusion/models/dream/configuration.json

@@ -1,14 +0,0 @@
-{
-  "framework": "pytorch",
-  "task" : "auto-speech-recognition",
-  "model": {"type" : "funasr"},
-  "pipeline": {"type":"funasr-pipeline"},
-  "model_name_in_hub": {
-    "ms":"", 
-    "hf":""},
-  "file_path_metas": {
-    "config":"config.yaml",
-    "tokenizer_conf": {"bpemodel": "chn_jpn_yue_eng_ko_spectok.bpe.model"},
-    "frontend_conf":{"cmvn_file": "am.mvn"}}
-}
-

omni_diffusion/models/dream/configuration_dream.py

@@ -1,86 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Dream model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_rope_utils import rope_config_validation
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class DreamConfig(PretrainedConfig):
-    model_type = "Dream"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size=151936,
-        hidden_size=4096,
-        intermediate_size=22016,
-        num_hidden_layers=32,
-        num_attention_heads=32,
-        num_key_value_heads=32,
-        hidden_act="silu",
-        max_position_embeddings=32768,
-        initializer_range=0.02,
-        rms_norm_eps=1e-6,
-        use_cache=False,  # cache not used in diffusion
-        tie_word_embeddings=False,
-        rope_theta=10000.0,
-        rope_scaling=None,
-        use_sliding_window=False,
-        sliding_window=4096,
-        max_window_layers=28,
-        attention_dropout=0.0,
-        mask_token_id=151666,
-        pad_token_id=151643,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.use_sliding_window = use_sliding_window
-        self.sliding_window = sliding_window if use_sliding_window else None
-        self.max_window_layers = max_window_layers
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_dropout = attention_dropout
-        # Validate the correctness of rotary position embeddings parameters
-        # BC: if there is a 'type' field, move it to 'rope_type'.
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-        rope_config_validation(self)
-        
-        super().__init__(
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-        self.mask_token_id = mask_token_id
-        self.pad_token_id = pad_token_id

omni_diffusion/models/dream/generation_utils.py

@@ -1,623 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import warnings
-import copy
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
-
-import torch
-import torch.distributions as dists
-from torch.nn import functional as F
-from transformers import __version__
-from transformers.generation.configuration_utils import (
-    GenerationConfig
-)
-from transformers.utils import (
-    ModelOutput,
-    is_torchdynamo_compiling,
-    logging,
-)
-logger = logging.get_logger(__name__)
-from tqdm import tqdm
-
-
-def top_p_logits(logits, top_p=None):
-    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-    sorted_indices_to_remove = cumulative_probs > top_p
-    # Shift the indices to the right to keep the first token above the threshold
-    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-    sorted_indices_to_remove[..., 0] = 0
-
-    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
-    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
-    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
-    return logits
-
-def top_k_logits(logits, top_k=None):
-    top_k = min(top_k, logits.size(-1))  # Safety check
-    # Remove all tokens with a probability less than the last token of the top-k
-    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
-    return logits
-
-
-def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False, 
-                  repeat_penalty=1.0, max_position_penalty=1.0, past_x=None, mask_id=None,):
-    if temperature > 0:
-        logits = logits / temperature
-
-    if top_p is not None and top_p < 1:
-        logits = top_p_logits(logits, top_p)
-    if top_k is not None:
-        logits = top_k_logits(logits, top_k)
-
-    if repeat_penalty != 1.0:
-        select_mask = torch.logical_and((past_x != 0), (past_x != mask_id))
-        generated_tokens = set(past_x[select_mask].tolist())
-        for token in set(generated_tokens):
-            logits[:, token][logits[:, token] < 0] *= repeat_penalty
-            logits[:, token][logits[:, token] >= 0] /= repeat_penalty
-
-    if max_position_penalty != 1.0:
-        token_length = logits.shape[-2]
-        if token_length > 100:
-            penalty_map = [i / (token_length - 100) * (max_position_penalty - 1.0) + 1.0 
-                           for i in range(token_length - 100)]
-            penalty_map = torch.tensor(penalty_map).unsqueeze(-1).to(logits.device).to(logits.dtype)
-            penalty_map = torch.cat([torch.ones_like(logits[:100, :1]), penalty_map], dim=0)
-            penalty_map = penalty_map.repeat(1, logits.shape[-1])
-        
-            logits[logits < 0] *= penalty_map[logits < 0]
-            logits[logits >= 0] /= penalty_map[logits >= 0]
-
-    probs = torch.softmax(logits, dim=-1)
-    if temperature > 0:
-        try:
-            x0 = dists.Categorical(probs=probs).sample()
-            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
-        except:
-            confidence, x0 = probs.max(dim=-1)
-    else:
-        confidence, x0 = probs.max(dim=-1)
-
-    if margin_confidence:
-        sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
-        top1_probs = sorted_probs[:, 0]
-        top2_probs = sorted_probs[:, 1]
-        confidence = top1_probs - top2_probs
-
-    if neg_entropy:
-        epsilon = 1e-10
-        log_probs = torch.log(probs + epsilon)
-        confidence = torch.sum(probs * log_probs, dim=-1)
-
-    return confidence, x0
-
-
-@dataclass
-class DreamModelOutput(ModelOutput):
-    sequences: torch.LongTensor = None
-    history: Optional[Tuple[torch.FloatTensor]] = None
-
-
-class DreamGenerationConfig(GenerationConfig):
-    def __init__(self, **kwargs):
-        self.temperature: float = kwargs.pop("temperature", 0.0)
-        self.top_p: Optional[float] = kwargs.pop("top_p", None)
-        self.top_k: Optional[int] = kwargs.pop("top_k", None)
-        self.max_length = kwargs.pop("max_length", 20)
-        self.max_new_tokens = kwargs.pop("max_new_tokens", None)
-        # diffusion specific params
-        self.eps: float = kwargs.pop("eps", 1e-3)
-        self.steps: int = kwargs.pop("steps", 512)
-        self.alg: str = kwargs.pop("alg", 'origin')
-        self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)
-
-        # Parameters that define the output variables of `generate`
-        self.num_return_sequences: int = kwargs.pop("num_return_sequences", 1)
-        self.return_dict_in_generate: bool = kwargs.pop("return_dict_in_generate", False)
-        self.output_history: bool = kwargs.pop("output_history", False)
-
-        # Special tokens that can be used at generation time
-        self.mask_token_id = kwargs.pop("mask_token_id", None)
-        self.pad_token_id = kwargs.pop("pad_token_id", None)
-        self.bos_token_id = kwargs.pop("bos_token_id", None)
-        self.eos_token_id = kwargs.pop("eos_token_id", None)
-
-        # Wild card
-        self.generation_kwargs = kwargs.pop("generation_kwargs", {})
-
-        # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub
-        # interface.
-        self._from_model_config = kwargs.pop("_from_model_config", False)
-        self._commit_hash = kwargs.pop("_commit_hash", None)
-        self.transformers_version = kwargs.pop("transformers_version", __version__)
-
-        # Additional attributes without default values
-        if not self._from_model_config:
-            # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a
-            # model's default configuration file
-            for key, value in kwargs.items():
-                try:
-                    setattr(self, key, value)
-                except AttributeError as err:
-                    logger.error(f"Can't set {key} with value {value} for {self}")
-                    raise err
-
-        # Validate the values of the attributes
-        self.validate(is_init=True)
-
-    def validate(self, is_init=False):
-        pass
-
-class DreamGenerationMixin:
-    @staticmethod
-    def _expand_inputs_for_generation(
-        expand_size: int = 1,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.LongTensor] = None
-    ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
-        """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
-        # Do not call torch.repeat_interleave if expand_size is 1 because it clones
-        # the input tensor and thus requires more memory although no change is applied
-        if expand_size == 1:
-            return input_ids, attention_mask
-        if input_ids is not None:
-            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
-        if attention_mask is not None:
-            attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)
-        return input_ids, attention_mask
-
-    def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
-        """Performs validation related to the resulting generated length"""
-
-        # Can't throw warnings/exceptions during compilation
-        if is_torchdynamo_compiling():
-            return
-
-        # 1. Max length warnings related to poor parameterization
-        if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
-            # 20 is the default max_length of the generation config
-            warnings.warn(
-                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
-                "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
-                "generation.",
-                UserWarning,
-            )
-        if input_ids_length >= generation_config.max_length:
-            input_ids_string = "input_ids"
-            raise ValueError(
-                f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
-                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
-                " increasing `max_length` or, better yet, setting `max_new_tokens`."
-            )
-
-    def _prepare_generated_length(
-        self,
-        generation_config,
-        has_default_max_length,
-        input_ids_length,
-    ):
-        """Prepared max and min length in generation configs to avoid clashes between similar attributes"""
-
-        if generation_config.max_new_tokens is not None:
-            if not has_default_max_length and generation_config.max_length is not None:
-                logger.warning(
-                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
-                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
-                    "Please refer to the documentation for more information. "
-                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
-                )
-            generation_config.max_length = generation_config.max_new_tokens + input_ids_length
-
-        elif has_default_max_length:
-            if generation_config.max_length == DreamGenerationConfig().max_length:
-                generation_config.max_length = generation_config.max_length + input_ids_length
-                max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
-                if max_position_embeddings is not None:
-                    generation_config.max_length = min(generation_config.max_length, max_position_embeddings)
-
-        return generation_config
-
-    def _prepare_generation_config(
-        self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict
-    ) -> DreamGenerationConfig:
-        """
-        Prepares the base generation config, then applies any generation configuration options from kwargs. This
-        function handles retrocompatibility with respect to configuration files.
-        """
-        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
-        using_model_generation_config = False
-        if generation_config is None:
-            generation_config = DreamGenerationConfig.from_model_config(self.config)
-            using_model_generation_config = True
-
-        # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config`
-        # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an
-        # exception will be raised in `_validate_model_kwargs`
-        if not is_torchdynamo_compiling():
-            generation_config = copy.deepcopy(generation_config)
-            _kwargs = generation_config.update(**kwargs)
-            # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model
-            if not using_model_generation_config:
-                if generation_config.bos_token_id is None:
-                    generation_config.bos_token_id = self.generation_config.bos_token_id
-                if generation_config.eos_token_id is None:
-                    generation_config.eos_token_id = self.generation_config.eos_token_id
-                if generation_config.pad_token_id is None:
-                    generation_config.pad_token_id = self.generation_config.pad_token_id
-                if generation_config.mask_token_id is None:
-                    generation_config.mask_token_id = self.generation_config.mask_token_id
-
-        return generation_config
-
-    def _prepare_special_tokens(
-        self,
-        generation_config: DreamGenerationConfig,
-        device: Optional[Union[torch.device, str]] = None,
-    ):
-        """
-        Prepares the special tokens for generation, overwriting the generation config with their processed versions
-        converted to tensor.
-
-        Note that `generation_config` is changed in place and stops being serializable after this method is called.
-        That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the
-        function). However, if called outside `generate`, consider creating a copy of `generation_config` first.
-        """
-
-        # Convert special tokens to tensors
-        def _tensor_or_none(token, device=None):
-            if token is None:
-                return token
-
-            device = device if device is not None else self.device
-            if isinstance(token, torch.Tensor):
-                return token.to(device)
-            return torch.tensor(token, device=device, dtype=torch.long)
-
-        bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
-        eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
-        pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
-        mask_token_tensor = _tensor_or_none(generation_config.mask_token_id, device=device)
-
-        # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
-        if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
-            eos_token_tensor = eos_token_tensor.unsqueeze(0)
-
-        # Set pad token if unset (and there are conditions to do so)
-        if pad_token_tensor is None and eos_token_tensor is not None:
-            pad_token_tensor = eos_token_tensor[0]
-            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")
-
-        # Update generation config with the updated special tokens tensors
-        # NOTE: this must be written into a different attribute name than the one holding the original special tokens
-        # (in their non-tensor form), in order to enable end-to-end compilation. See
-        # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations
-        generation_config._bos_token_tensor = bos_token_tensor
-        generation_config._eos_token_tensor = eos_token_tensor
-        generation_config._pad_token_tensor = pad_token_tensor
-        generation_config._mask_token_tensor = mask_token_tensor
-
-    @torch.no_grad()
-    def diffusion_generate(
-        self,
-        inputs: Optional[torch.Tensor] = None,
-        generation_config: Optional[DreamGenerationConfig] = None,
-        inputs_embeds=None,
-        prefix_lm=False,
-        alg=None,
-        block_size=-1,
-        cfg=0.0,
-        add_boa_token=False,
-        **kwargs,
-    ) -> Union[DreamModelOutput, torch.LongTensor]:
-        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
-        generation_config = self._prepare_generation_config(generation_config, **kwargs)
-        generation_tokens_hook_func = kwargs.pop("generation_tokens_hook_func", lambda step, x, logits: x)
-        generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)
-        # breakpoint()
-        # 2. Define model inputs
-        if inputs is not None:
-            input_ids = inputs
-            device = input_ids.device
-            input_ids_length = input_ids.shape[-1]
-        else:
-            input_ids = None
-            device = inputs_embeds.device
-            input_ids_length = inputs_embeds.shape[1]
-        attention_mask = kwargs.pop("attention_mask", None)
-        self._prepare_special_tokens(generation_config, device=device)
-
-        # 3. Prepare `max_length`.
-        
-        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
-        generation_config = self._prepare_generated_length(
-            generation_config=generation_config,
-            has_default_max_length=has_default_max_length,
-            input_ids_length=input_ids_length,
-        )
-
-        self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
-        # import pdb;pdb.set_trace()
-        # 4. Check input_ids
-        #if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
-        if not is_torchdynamo_compiling() and self.device.type != device.type:
-            warnings.warn(
-                "You are calling .generate() with the `input_ids` being on a device type different"
-                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
-                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
-                " Please make sure that you have put `input_ids` to the"
-                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
-                " running `.generate()`.",
-                UserWarning,
-            )
-        # breakpoint()
-        if (
-            hasattr(generation_config, "pad_token_id") and
-            input_ids is not None and
-            torch.any(input_ids == generation_config.pad_token_id) and 
-            attention_mask is None
-        ):
-            warnings.warn(
-                "Padding was detected but no attention mask is passed here. For correct "
-                "generation results, please set `attention_mask` when batch-padding inputs.",
-                UserWarning,
-            )
-        assert generation_config.num_return_sequences == 1, \
-            "Currently, we only support num_return_sequences = 1 for diffusion generation."
-        input_ids, attention_mask = self._expand_inputs_for_generation(
-            expand_size=generation_config.num_return_sequences,
-            input_ids=input_ids,
-            attention_mask=attention_mask 
-        )
-        result = self._sample(
-            input_ids,
-            attention_mask=attention_mask,
-            generation_config=generation_config,
-            generation_tokens_hook_func=generation_tokens_hook_func,
-            generation_logits_hook_func=generation_logits_hook_func,
-            inputs_embeds=inputs_embeds,
-            device=device,
-            prefix_lm=prefix_lm,
-            alg=alg,
-            block_size=block_size,
-            cfg=cfg,
-            add_boa_token=add_boa_token,
-            **kwargs,
-        )
-        return result
-
-    def _sample(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor],
-        generation_config: DreamGenerationConfig,
-        generation_tokens_hook_func,
-        generation_logits_hook_func,
-        inputs_embeds=None,
-        prefix_lm=False,
-        device=None,
-        step_ratio=None,
-        penalty=1.2,
-        alg=None,
-        block_size=None,
-        add_boa_token=False,
-        max_position_penalty=1.0,
-        repeat_penalty=1.0,
-        cfg=0.0,
-        **kwargs,
-    ) -> Union[DreamModelOutput, torch.LongTensor]:
-        output_history = True
-        return_dict_in_generate = generation_config.return_dict_in_generate  
-        max_length = generation_config.max_length                     
-        mask_token_id = generation_config.mask_token_id               
-        max_new_tokens = generation_config.max_new_tokens             
-        steps = min(generation_config.steps, max_new_tokens) 
-        eps = generation_config.eps                                    
-        alg = generation_config.alg if alg is None else alg                       
-        print("denoise algorithm: " + alg)
-        alg_temp = generation_config.alg_temp                          
-        temperature = generation_config.temperature                    
-        top_p = generation_config.top_p                                
-        top_k = generation_config.top_k              
-
-        histories = [] if (return_dict_in_generate and output_history) else None
-        all_logit = []
-        generated_tokens = []
-        block_size = max_new_tokens if block_size < 0 else block_size
-
-        if input_ids is None:
-            assert device is not None
-            assert inputs_embeds is not None
-            bsz, seq_len = inputs_embeds.shape[:2]                  
-            max_length = seq_len + max_new_tokens                   
-            input_ids = torch.full((bsz, seq_len), 0, dtype=torch.long).to(device)
-
-        tok_idx = None
-        past_key_values = None
-
-        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id) 
-
-        timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
-        x = generation_tokens_hook_func(None, x, None)
-
-        if step_ratio is not None:
-            steps = int(max_new_tokens * step_ratio)
-
-        if add_boa_token:
-            bos_index = int((x.shape[1] - (x == mask_token_id).sum()) + (x == mask_token_id).sum() * 0.2)
-            x[:, bos_index] = 151684 # <|begin_of_audio|>
-
-        input_x = x.clone()
-        total_steps = steps
-        block_num = (x == mask_token_id).sum() // block_size
-        if block_num * block_size < (x == mask_token_id).sum(): block_num += 1
-        input_length = input_ids.shape[-1]
-
-        task = None
-        if "task" in kwargs: task = kwargs['task']
-        if cfg > 0:
-            import random
-            empty_prompt = ""
-            if task == "S2I":
-                empty_prompt = "<|im_start|>system\nPlease generate an image based on the input audio.<|im_end|>\n"
-                empty_prompt += "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n"
-                un_x = kwargs['tokenizer'].encode(empty_prompt)
-            elif task == "T2I":
-                empty_prompt = "<|im_start|>user\nGenerate an image based on the provided text description.\n"
-                empty_prompt += "<|im_end|>\n<|im_start|>assistant\n"
-                first_audio_token = kwargs['tokenizer'].encode("<|begin_of_audio|>")[0]
-                un_x_text = random.sample([_ for _ in range(first_audio_token)], 
-                                          input_ids.shape[1] - len(kwargs['tokenizer'].encode(empty_prompt)))
-                un_x = kwargs['tokenizer'].encode("<|im_start|>user\nGenerate an image based on the provided \
-                                                  text description.\n") 
-                un_x = un_x + un_x_text + kwargs['tokenizer'].encode("<|im_end|>\n<|im_start|>assistant\n")
-
-        for block_idx in range(block_num):
-            block_mask = torch.zeros([x.shape[-1]]).to(torch.bool).to(x.device)
-            block_mask[input_length + block_idx * block_size: input_length + (block_idx + 1) * block_size] = True
-            steps = int(block_mask.sum() / (x.shape[-1] - input_length) * total_steps)
-            timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
-            for i in tqdm(range(steps)):
-                mask_index = (x == mask_token_id)
-                if mask_index.sum() == 0: break
-                inputs_embeds_curr = self.model.embed_tokens(x)
-
-                if inputs_embeds is not None:
-                    inputs_embeds_curr[:, :inputs_embeds.shape[1]] = inputs_embeds
-
-                if cfg > 0:
-                    input_un_x = torch.tensor(un_x).unsqueeze(0).to(x.dtype).to(x.device)
-                    input_un_x = torch.cat([input_un_x, x[:, input_ids.shape[1]:]], dim=1)
-                    un_inpus_embeds = self.model.embed_tokens(input_un_x)
-
-                    attention_mask_cond = torch.ones([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
-                    attention_mask_cond = attention_mask_cond.to(torch.bool).to(inputs_embeds_curr.device)
-                    attention_mask_uncond = torch.zeros([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
-                    attention_mask_uncond[:, :un_inpus_embeds.shape[1], :un_inpus_embeds.shape[1]] = 1
-                    attention_mask_uncond = attention_mask_uncond.to(torch.bool).to(inputs_embeds.device)
-                    attention_mask = torch.cat([attention_mask_cond, attention_mask_uncond])
-                    attention_mask = attention_mask.unsqueeze(1)
-
-                    if inputs_embeds_curr.shape[1] != un_inpus_embeds.shape[1]:
-                        un_inpus_embeds = torch.cat([un_inpus_embeds, 
-                            torch.zeros_like(inputs_embeds_curr[:, :inputs_embeds_curr.shape[1] - 
-                                                                un_inpus_embeds.shape[1], :])], dim=1)
-                    input_inputs_embeds_curr = torch.cat([inputs_embeds_curr, un_inpus_embeds])
-
-                    model_logits = self.forward_dream(None, attention_mask, tok_idx, 
-                                                      inputs_embeds=input_inputs_embeds_curr).logits 
-                    logits = model_logits[:1]; un_logits = model_logits[1:]
-                    logits = un_logits + (cfg + 1) * (logits - un_logits)
-                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
-
-                else:
-                    logits = self.forward_dream(None, attention_mask, tok_idx, 
-                                                inputs_embeds=inputs_embeds_curr).logits
-                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
-
-                logits = generation_logits_hook_func(i, x, logits)
-
-                mask_logits = logits[mask_index]
-                if i == 0:
-                    input_index = torch.where(mask_index[0]==True)[0][0]
-
-                t = timesteps[i]
-                s = timesteps[i + 1]
-
-                if alg == 'origin':
-                    p_transfer = 1 - s / t if i < steps - 1 else 1  
-                    x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id
-                    transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer
-                    _, x0[transfer_index_t_s] = sample_tokens(
-                        mask_logits[transfer_index_t_s],
-                        temperature=temperature,
-                        top_p=top_p,
-                        top_k=top_k,
-                        max_position_penalty=max_position_penalty,
-                    )
-                    x[mask_index] = x0.clone()
-
-                else:
-                    if alg == 'maskgit_plus':
-                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, 
-                                                       max_position_penalty=max_position_penalty)
-                    elif alg == 'topk_margin':
-                        confidence, x0 = sample_tokens(
-                            mask_logits,
-                            temperature=temperature,
-                            top_p=top_p,
-                            top_k=top_k,
-                            margin_confidence=True,
-                            max_position_penalty=max_position_penalty,
-                        )
-                    elif alg == 'entropy':
-                        confidence, x0 = sample_tokens(
-                            mask_logits,
-                            temperature,
-                            top_p=top_p,
-                            top_k=top_k,
-                            neg_entropy=True,
-                            max_position_penalty=max_position_penalty,
-                        )
-                    elif alg == "entropy-penalty": 
-                        confidence, x0 = sample_tokens(
-                            mask_logits,
-                            temperature,
-                            top_p=top_p,
-                            top_k=top_k,
-                            neg_entropy=True,
-                            repeat_penalty=repeat_penalty if len(histories) != 0 else 1.0,
-                            past_x=histories[-1] if len(histories) != 0 else [],
-                            mask_id=mask_token_id,
-                            max_position_penalty=max_position_penalty,
-                        )
-                    else:
-                        raise RuntimeError(f"Unknown alg: {alg}")
-
-                    block_mask_1 = block_mask[mask_index[0]]
-                    confidence = confidence + torch.where(block_mask_1, 0, -torch.inf).to(confidence.device)
-
-                    num_mask_token = mask_index.sum()
-                    num_mask_token = (x[:, block_mask] == mask_token_id).sum()
-                    number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else num_mask_token
-                    number_transfer_tokens = max(number_transfer_tokens, 1)
-
-                    if number_transfer_tokens > 0:
-                        if alg_temp is None or alg_temp == 0:
-                            _, transfer_index = torch.topk(confidence, number_transfer_tokens)
-                        else:
-                            confidence = confidence / alg_temp
-                            confidence = F.softmax(confidence, dim=-1)
-                            transfer_index = torch.multinomial(confidence, num_samples=number_transfer_tokens)
-
-                        x0_ = torch.zeros_like(x0, device=self.device, dtype=torch.long) + mask_token_id
-                        x0_[transfer_index] = x0[transfer_index].clone()
-                        x[mask_index] = x0_
-
-                        logit,indic = torch.max(torch.softmax(logits.clone(),dim=-1),-1)
-                        logit = logit[0][x[0]!=0]
-                        indic = indic[0][x[0]!=0]
-                        temp_X = x[0][x[0]!=0]
-
-                x = generation_tokens_hook_func(i, x, logits)
-
-                if histories is not None:
-                    histories.append(x.clone())
-                    all_logit.append(torch.max(logits.clone(),-1)[-1])
-
-        return (x, histories)

omni_diffusion/models/dream/modeling_dream.py

@@ -1,1136 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT and Qwen implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT and Qwen used by the Meta AI and Qwen team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch Dream model."""
-
-
-
-from .modeling_sensevoice import AudioEncoder
-from .resampler_projector import ResamplerProjector
-import random
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-
-import math
-from typing import List, Optional, Tuple, Union
-import os
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_outputs import (
-    # BaseModelOutput,
-    MaskedLMOutput,
-    ModelOutput
-)
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import (
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    is_flash_attn_2_available,
-    is_flash_attn_greater_or_equal_2_10,
-    logging,
-)
-from transformers import PretrainedConfig
-from .configuration_dream import DreamConfig
-from .generation_utils import DreamGenerationMixin, DreamGenerationConfig
-from  dataclasses import  dataclass
-from typing import Any
-if is_flash_attn_2_available():
-    from transformers.modeling_flash_attention_utils import _flash_attention_forward
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-from .modeling_sensevoice import AudioEncoder
-from .resampler_projector import ResamplerProjector
-
-import torch
-import random
-
-
-logger = logging.get_logger(__name__)
-
-
-@dataclass
-class MaskedLMOutput(ModelOutput):
-    """
-    Base class for masked language models outputs.
-
-    Args:
-        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
-            Masked language modeling (MLM) loss.
-        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
-            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
-        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed
-        or when `config.output_hidden_states=True`):
-                Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has 
-                an embedding layer, + one for the output of each layer) of shape `(batch_size, 
-                sequence_length, hidden_size)`.
-
-            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
-        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed 
-        or when `config.output_attentions=True`):
-            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
-            sequence_length)`.
-
-            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-            heads.
-    """
-
-    loss: Optional[torch.FloatTensor] = None
-    logits: torch.FloatTensor = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
-
-_CHECKPOINT_FOR_DOC = "Dream-7B"
-_CONFIG_FOR_DOC = "DreamConfig"
-import os
-ENFORCE_NUM_ITEMIN_BATCH = os.environ.get("ENFORCE_NUM_ITEMIN_BATCH", False)
-
-@dataclass
-class BaseModelOutput(ModelOutput):
-    """
-    Base class for model's outputs, with potential hidden states and attentions.
-
-    Args:
-        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            Sequence of hidden-states at the output of the last layer of the model.
-        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed 
-        or when `config.output_hidden_states=True`):
-            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
-            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
-            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
-        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed 
-        or when `config.output_attentions=True`):
-            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
-            sequence_length)`.
-
-            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-            heads.
-    """
-
-    last_hidden_state: torch.FloatTensor = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    past_key_values: Optional[Cache] = None
-# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Dream
-
-
-class DreamRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        """
-        DreamRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Dream
-class DreamRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim=None,
-        max_position_embeddings=2048,
-        base=10000,
-        device=None,
-        scaling_factor=1.0,
-        rope_type="default",
-        config: Optional[DreamConfig] = None,
-    ):
-        super().__init__()
-        # TODO (joao): remove the `if` below, only used for BC
-        self.rope_kwargs = {}
-        if config is None:
-            logger.warning_once(
-                "`DreamRotaryEmbedding` can now be fully parameterized by passing the model config through the "
-                "`config` argument. All other arguments will be removed in v4.46"
-            )
-            self.rope_kwargs = {
-                "rope_type": rope_type,
-                "factor": scaling_factor,
-                "dim": dim,
-                "base": base,
-                "max_position_embeddings": max_position_embeddings,
-            }
-            self.rope_type = rope_type
-            self.max_seq_len_cached = max_position_embeddings
-            self.original_max_seq_len = max_position_embeddings
-        else:
-            # BC: "rope_type" was originally "type"
-            if config.rope_scaling is not None:
-                self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
-            else:
-                self.rope_type = "default"
-            self.max_seq_len_cached = config.max_position_embeddings
-            self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    def reset_parameters(self):
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, self.inv_freq.device, **self.rope_kwargs)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-    
-
-    def _dynamic_frequency_update(self, position_ids, device):
-        """
-        dynamic RoPE layers should recompute `inv_freq` in the following situations:
-        1 - growing beyond the cached sequence length (allow scaling)
-        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
-        """
-        seq_len = torch.max(position_ids) + 1
-        if seq_len > self.max_seq_len_cached:  # growth
-            inv_freq, self.attention_scaling = self.rope_init_fn(
-                self.config, device, seq_len=seq_len, **self.rope_kwargs
-            )
-            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
-            self.max_seq_len_cached = seq_len
-
-        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
-            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
-            self.max_seq_len_cached = self.original_max_seq_len
-
-    @torch.no_grad()
-    def forward(self, x, position_ids):
-        if "dynamic" in self.rope_type:
-            self._dynamic_frequency_update(position_ids, device=x.device)
-
-        # Core RoPE block
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
-        device_type = x.device.type
-        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos()
-            sin = emb.sin()
-
-        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
-        cos = cos * self.attention_scaling
-        sin = sin * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-# Copied from transformers.models.llama.modeling_llama.rotate_half
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Dream
-class DreamMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, hidden_state):
-        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
-
-
-# Copied from transformers.models.llama.modeling_llama.repeat_kv
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-class DreamAttention(nn.Module):
-    """
-    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
-    and "Generating Long Sequences with Sparse Transformers".
-    """
-
-    def __init__(self, config: DreamConfig, layer_idx: Optional[int] = None):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        if layer_idx is None:
-            logger.warning_once(
-                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
-                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
-                "when creating this class."
-            )
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = self.hidden_size // self.num_heads
-        self.num_key_value_heads = config.num_key_value_heads
-        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-        self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
-        self.is_causal = False
-        self.attention_dropout = config.attention_dropout
-
-        if (self.head_dim * self.num_heads) != self.hidden_size:
-            raise ValueError(
-                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
-                f" and `num_heads`: {self.num_heads})."
-            )
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
-
-        self.rotary_emb = DreamRotaryEmbedding(config=self.config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
-        output_attentions: bool = False,
-        use_cache: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        if position_embeddings is None:
-            logger.warning_once(
-                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
-                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
-                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
-                "removed and `position_embeddings` will be mandatory."
-            )
-            cos, sin = self.rotary_emb(value_states, position_ids)
-        else:
-            cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        # repeat k/v heads if n_kv_heads < n_heads
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-        if attention_mask is not None:  # no matter the length, we just slice it
-            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-            attn_weights = attn_weights + causal_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
-        attn_output = self.o_proj(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, past_key_value
-    
-
-class DreamSdpaAttention(DreamAttention):
-    """
-    Dream attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `DreamAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
-    SDPA API.
-    """
-
-    # Adapted from DreamAttention.forward
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
-        use_cache: bool = False,
-        output_attentions: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        if output_attentions:
-            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` 
-            # once this is implemented.
-            logger.warning_once(
-                "DreamModel is using DreamSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` " \
-                " does not support" \
-                " `output_attentions=True`. Falling back to the manual attention implementation, "
-                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. ' \
-                'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
-            )
-            return super().forward(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-            )
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        if position_embeddings is None:
-            logger.warning_once(
-                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
-                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
-                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
-                "removed and `position_embeddings` will be mandatory."
-            )
-            cos, sin = self.rotary_emb(value_states, position_ids)
-        else:
-            cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs 
-        # with custom attn_mask,
-        # Reference: https://github.com/pytorch/pytorch/issues/112577.
-        if query_states.device.type == "cuda" and attention_mask is not None:
-            query_states = query_states.contiguous()
-            key_states = key_states.contiguous()
-            value_states = value_states.contiguous()
-
-        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement 
-        # instead of an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes 
-        # and full graph options. An inline conditional prevents dynamic shapes from compiling.
-        # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create 
-        # a causal mask in case q_len == 1.
-
-        if attention_mask == None:
-            attention_mask = torch.ones([1, 1, 
-            query_states.shape[2], query_states.shape[2]]).to(torch.bool).to(query_states.device)
-        bool_mask = attention_mask.to(torch.bool)
-
-        attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=bool_mask ,
-            dropout_p=self.attention_dropout if self.training else 0.0,
-            is_causal=False, # hard coded
-        )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(bsz, q_len, self.hidden_size)
-
-        attn_output = self.o_proj(attn_output)
-
-        return attn_output, None, past_key_value
-
-class DreamDecoderLayer(nn.Module):
-    def __init__(self, config: DreamConfig, layer_idx: int):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        if config.sliding_window and config._attn_implementation != "flash_attention_2":
-            logger.warning_once(
-                f"Sliding Window Attention is enabled but not implemented for "
-                f"`{config._attn_implementation}`; "
-                "unexpected results may be encountered."
-            )
-        
-        # self.self_attn = Dream_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
-        self.self_attn = DreamSdpaAttention(config, layer_idx)
-
-        self.mlp = DreamMLP(config)
-        self.input_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    # @torch.compile
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, sequence_length)` where padding elements are indicated by 0.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-                Indices depicting the position of the input sequence tokens in the sequence.
-            position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
-                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
-                with `head_dim` being the embedding dimension of each attention head.
-            kwargs (`dict`, *optional*):
-                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
-                into the model
-        """
-
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            position_embeddings=position_embeddings,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        if use_cache:
-            outputs += (present_key_value,)
-
-        return outputs
-
-class DreamPreTrainedModel(PreTrainedModel):
-    config_class = DreamConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["DreamDecoderLayer"]
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
-    _supports_cache_class = True
-    _supports_quantized_cache = True
-    _supports_static_cache = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
-        *model_args,
-        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
-        cache_dir: Optional[Union[str, os.PathLike]] = None,
-        ignore_mismatched_sizes: bool = False,
-        force_download: bool = False,
-        local_files_only: bool = False,
-        token: Optional[Union[str, bool]] = None,
-        revision: str = "main",
-        use_safetensors: Optional[bool] = None,
-        weights_only: bool = True,
-        **kwargs,
-    ):
-        _ = None
-        try:
-            _model,_  = super().from_pretrained(
-                pretrained_model_name_or_path,
-                *model_args,
-                config=config,
-                cache_dir=cache_dir,
-                ignore_mismatched_sizes=ignore_mismatched_sizes,
-                force_download=force_download,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                use_safetensors=use_safetensors,
-                weights_only=weights_only,
-                **kwargs,
-            )
-        except Exception as e:
-            _model = super().from_pretrained(
-                pretrained_model_name_or_path,
-                *model_args,
-                config=config,
-                cache_dir=cache_dir,
-                ignore_mismatched_sizes=ignore_mismatched_sizes,
-                force_download=force_download,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                use_safetensors=use_safetensors,
-                weights_only=weights_only,
-                **kwargs,
-            )
-
-        resume_download = kwargs.get("resume_download", None)
-        proxies = kwargs.get("proxies", None)
-        subfolder = kwargs.get("subfolder", "")
-        from_auto_class = kwargs.get("_from_auto", False)
-        from_pipeline = kwargs.get("_from_pipeline", None)
-        _model.generation_config= DreamGenerationConfig.from_pretrained(
-            pretrained_model_name_or_path,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            subfolder=subfolder,
-            _from_auto=from_auto_class,
-            _from_pipeline=from_pipeline,
-        )
-        if _ is not None:
-            return _model,_
-        return _model
-    
-class DreamPrefixLMCache(Cache):
-    
-    def __init__(self):
-        super().__init__()
-        self.past_key_values = {}
-        # this will not be updated beyond the prefilling phase
-        
-    def update(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        layer_idx: int,
-        cache_kwargs = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if layer_idx in self.past_key_values:
-            past_key, past_value = self.past_key_values[layer_idx]
-            key_states = torch.cat((past_key, key_states), dim=-2)
-            value_states = torch.cat((past_value, value_states), dim=-2)
-            return key_states,value_states
-        else:
-            self.past_key_values[layer_idx] = (key_states, value_states)
-            return key_states, value_states
-
-    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
-        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
-        # TODO: deprecate this function in favor of `cache_position`
-        if len(self.past_key_values) == 0:
-            return 0
-        else:
-            return self.past_key_values[0][0].shape[-2]
-        
-    def get_max_cache_shape(self) -> Optional[int]:
-        return None
- 
-
-import deepspeed
-class DreamBaseModel(DreamPreTrainedModel):#
-    """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DreamDecoderLayer`]
-
-    Args:
-        config: DreamConfig
-    """
-
-    def __init__(self, config: DreamConfig):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = nn.ModuleList(
-            [DreamDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self._attn_implementation = config._attn_implementation
-        self.norm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = DreamRotaryEmbedding(config=config)
-
-        self.gradient_checkpointing = False
-        # Initialize weights and apply final processing
-
-        self.audio_model = AudioEncoder()
-        self.audio_projection = ResamplerProjector(512, config.hidden_size)
-
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        audios: Optional[torch.FloatTensor] = None,
-        audio_indices: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-    ) -> Union[Tuple, BaseModelOutput]:
-
-        if (past_key_values is None or len(past_key_values) == 0) and audios is not None:
-            audio_embeds, audio_lengths = self.audio_model(audios)
-            assert audio_embeds.shape[0] == len(audios)
-            fake_audios = None
-
-            audio_embeds = self.audio_projection(audio_embeds)
-
-        elif self.training:
-            device = self.get_input_embeddings().weight.data.device
-            dtype = self.get_input_embeddings().weight.data.dtype
-            fake_audios = torch.ones((1, 1, 560), dtype=dtype, device=device)
-            audio_embeds, audio_lengths = self.audio_model(fake_audios)
-            audio_embeds = self.audio_projection(audio_embeds)
-
-        else:
-            fake_audios = None
-            audio_embeds = None
-
-    
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-        
-
-
-        if fake_audios is not None:
-            inputs_embeds = inputs_embeds + audio_embeds.mean() * 0.0
-        elif audio_embeds is not None:
-            inputs_embeds = inputs_embeds.clone()
-            for audio_embeds_, audio_lengths_, audio_indices_ in zip(audio_embeds, audio_lengths, audio_indices,):
-                # print(f"{audio_embeds_.size()=} {audio_lengths_=} {audio_indices_.size()=}")
-                audio_embeds_ = audio_embeds_[:audio_lengths_, ...]
-                audio_embeds_ = audio_embeds_.to(inputs_embeds.device)
-                indices_b, indices_s = audio_indices_.to(inputs_embeds.device).unbind(dim=0)
-                inputs_embeds[indices_b.view(-1), indices_s.view(-1)] = audio_embeds_.view(-1, audio_embeds_.shape[-1])
-
-
-        if use_cache and past_key_values is None:
-            past_key_values = DreamPrefixLMCache()
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-
-        for decoder_layer in self.layers:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-                layer_outputs = deepspeed.checkpointing.checkpoint(
-                    decoder_layer,
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    past_key_values,
-                    output_attentions,
-                    use_cache,
-                    cache_position,
-                    position_embeddings,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_values,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                    position_embeddings=position_embeddings,
-                )
-
-            # breakpoint()
-            if isinstance(layer_outputs,torch.Tensor):
-                layer_outputs = (layer_outputs,None)
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutput(
-            last_hidden_state=hidden_states,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-            past_key_values=past_key_values,
-        )
-
-
-class DreamModel(DreamGenerationMixin, DreamPreTrainedModel):
-    _tied_weights_keys = ["lm_head.weight"]
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = DreamBaseModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-        # Initialize weights and apply final processing
-        self.tokenizer = None
-        self.post_init()
-
-    def reset_rope_parameters(self):
-        self.model.rotary_emb.reset_parameters()
-        for layer in self.model.layers:
-            layer.self_attn.rotary_emb.reset_parameters()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        audios: Optional[torch.FloatTensor] = None,
-        audio_indices: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        num_logits_to_keep: int = 0,
-        num_items_in_batch: int = None,
-        **loss_kwargs,
-    ) -> Union[Tuple, MaskedLMOutput]:
-
-        num_items_in_batch = None
-        if ENFORCE_NUM_ITEMIN_BATCH:
-            num_items_in_batch = labels.ne(-100).sum()
-            num_items_in_batch = torch.distributed.reduce(num_items_in_batch)
-
-        is_new = position_ids == 0
-        # is_new[0] = True
-        segment_id = torch.cumsum(is_new.long(), dim=1) - 1
-        new_attention_mask = (segment_id.unsqueeze(1) == segment_id.unsqueeze(2)).long()
-        mask = attention_mask.unsqueeze(-1)  # [bs, len, 1]
-        new_attention_mask = new_attention_mask * mask 
-
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=new_attention_mask,
-            audios=audios,
-            audio_indices=audio_indices,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            cache_position=cache_position,
-        )
-        hidden_states = outputs[0]
-        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
-
-        loss = None
-        if labels is not None:
-            if ENFORCE_NUM_ITEMIN_BATCH:
-                assert num_items_in_batch is not None, \
-                    "num_items_in_batch must be provided if ENFORCE_NUM_ITEMIN_BATCH is True"
-            loss = self.loss_function(logits, labels, 
-                                      self.vocab_size,num_items_in_batch=num_items_in_batch, **loss_kwargs)
-
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-
-        return MaskedLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-            past_key_values=outputs.past_key_values
-        )
-        
-    def forward_dream(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        num_logits_to_keep: int = 0,
-        **loss_kwargs,
-    ) -> Union[Tuple, MaskedLMOutput]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        attention_mask = None
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            cache_position=cache_position,
-        )
-
-        hidden_states = outputs[0]
-        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-
-        return MaskedLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-            past_key_values=outputs.past_key_values,
-        )
-
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        audios: Optional[torch.FloatTensor] = None,
-        audio_indices: Optional[torch.LongTensor] = None,
-        max_new_tokens=512,
-        steps=512,
-        temperature=0.2,
-        top_p=0.95,
-        alg_temp=0.,
-        alg="entropy",
-        output_history=False,
-        **kwargs,
-    ):
-        # modalities = kwargs.pop("modalities", None) if "modalities" in kwargs and modalities is None else modalities
-        position_ids = kwargs.pop("position_ids", None)
-        attention_mask = kwargs.pop("attention_mask", None)
-        if "inputs_embeds" in kwargs:
-            raise NotImplementedError("`inputs_embeds` is not supported")
-
-        if audios is not None:
-            audio_embeds, audio_lengths = self.model.audio_model(audios)
-            assert audio_embeds.shape[0] == len(audios)
-            audio_embeds = self.model.audio_projection(audio_embeds)
-        else:
-            audio_embeds = None
-
-        inputs_embeds = self.model.embed_tokens(input_ids)
-        
-        if audio_embeds is not None:
-            inputs_embeds = inputs_embeds.clone()
-            for audio_embeds_, audio_lengths_, audio_indices_ in zip(audio_embeds, audio_lengths, audio_indices,):
-                # print(f"{audio_embeds_.size()=} {audio_lengths_=} {audio_indices_.size()=}")
-                audio_embeds_ = audio_embeds_[:audio_lengths_, ...]
-                audio_embeds_ = audio_embeds_.to(inputs_embeds.device)
-                indices_b, indices_s = audio_indices_.to(inputs_embeds.device).unbind(dim=0)
-                inputs_embeds[indices_b.view(-1), indices_s.view(-1)] = audio_embeds_.view(-1, audio_embeds_.shape[-1])
-
-        return self.diffusion_generate(
-            None,
-            inputs_embeds=inputs_embeds,
-            max_new_tokens=max_new_tokens,
-            output_history=output_history,
-            return_dict_in_generate=True,
-            steps=steps,
-            temperature=temperature,
-            top_p=top_p,
-            alg=alg,
-            alg_temp=alg_temp,
-            **kwargs
-        )

omni_diffusion/models/dream/modeling_sensevoice.py

@@ -1,1258 +0,0 @@
-
-import time
-import torch
-from torch import nn
-import torch.nn.functional as F
-from typing import Iterable, Optional
-
-from funasr.register import tables
-from funasr.models.ctc.ctc import CTC
-from funasr.utils.datadir_writer import DatadirWriter
-from funasr.models.paraformer.search import Hypothesis
-from funasr.train_utils.device_funcs import force_gatherable
-from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
-from funasr.metrics.compute_acc import compute_accuracy, th_accuracy
-from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
-# from utils.ctc_alignment import ctc_forced_align
-
-def ctc_forced_align(
-    log_probs: torch.Tensor,
-    targets: torch.Tensor,
-    input_lengths: torch.Tensor,
-    target_lengths: torch.Tensor,
-    blank: int = 0,
-    ignore_id: int = -1,
-) -> torch.Tensor:
-    """Align a CTC label sequence to an emission.
-
-    Args:
-        log_probs (Tensor): log probability of CTC emission output.
-            Tensor of shape `(B, T, C)`. where `B` is the batch size, `T` is the input length,
-            `C` is the number of characters in alphabet including blank.
-        targets (Tensor): Target sequence. Tensor of shape `(B, L)`,
-            where `L` is the target length.
-        input_lengths (Tensor):
-            Lengths of the inputs (max value must each be <= `T`). 1-D Tensor of shape `(B,)`.
-        target_lengths (Tensor):
-            Lengths of the targets. 1-D Tensor of shape `(B,)`.
-        blank_id (int, optional): The index of blank symbol in CTC emission. (Default: 0)
-        ignore_id (int, optional): The index of ignore symbol in CTC emission. (Default: -1)
-    """
-    targets[targets == ignore_id] = blank
-
-    batch_size, input_time_size, _ = log_probs.size()
-    bsz_indices = torch.arange(batch_size, device=input_lengths.device)
-
-    _t_a_r_g_e_t_s_ = torch.cat(
-        (
-            torch.stack((torch.full_like(targets, blank), targets), dim=-1).flatten(start_dim=1),
-            torch.full_like(targets[:, :1], blank),
-        ),
-        dim=-1,
-    )
-    diff_labels = torch.cat(
-        (
-            torch.as_tensor([[False, False]], device=targets.device).expand(batch_size, -1),
-            _t_a_r_g_e_t_s_[:, 2:] != _t_a_r_g_e_t_s_[:, :-2],
-        ),
-        dim=1,
-    )
-
-    neg_inf = torch.tensor(float("-inf"), device=log_probs.device, dtype=log_probs.dtype)
-    padding_num = 2
-    padded_t = padding_num + _t_a_r_g_e_t_s_.size(-1)
-    best_score = torch.full((batch_size, padded_t), neg_inf, device=log_probs.device, dtype=log_probs.dtype)
-    best_score[:, padding_num + 0] = log_probs[:, 0, blank]
-    best_score[:, padding_num + 1] = log_probs[bsz_indices, 0, _t_a_r_g_e_t_s_[:, 1]]
-
-    backpointers = torch.zeros((batch_size, input_time_size, padded_t), device=log_probs.device, dtype=targets.dtype)
-
-    for t in range(1, input_time_size):
-        prev = torch.stack(
-            (best_score[:, 2:], best_score[:, 1:-1], torch.where(diff_labels, best_score[:, :-2], neg_inf))
-        )
-        prev_max_value, prev_max_idx = prev.max(dim=0)
-        best_score[:, padding_num:] = log_probs[:, t].gather(-1, _t_a_r_g_e_t_s_) + prev_max_value
-        backpointers[:, t, padding_num:] = prev_max_idx
-
-    l1l2 = best_score.gather(
-        -1, torch.stack((padding_num + target_lengths * 2 - 1, padding_num + target_lengths * 2), dim=-1)
-    )
-
-    path = torch.zeros((batch_size, input_time_size), device=best_score.device, dtype=torch.long)
-    path[bsz_indices, input_lengths - 1] = padding_num + target_lengths * 2 - 1 + l1l2.argmax(dim=-1)
-
-    for t in range(input_time_size - 1, 0, -1):
-        target_indices = path[:, t]
-        prev_max_idx = backpointers[bsz_indices, t, target_indices]
-        path[:, t - 1] += target_indices - prev_max_idx
-
-    alignments = _t_a_r_g_e_t_s_.gather(dim=-1, index=(path - padding_num).clamp(min=0))
-    return alignments
-
-class SinusoidalPositionEncoder(torch.nn.Module):
-    """ """
-
-    def __int__(self, d_model=80, dropout_rate=0.1):
-        pass
-
-    def encode(
-        self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32
-    ):
-        batch_size = positions.size(0)
-        positions = positions.type(dtype)
-        device = positions.device
-        log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / (
-            depth / 2 - 1
-        )
-        inv_timescales = torch.exp(
-            torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment)
-        )
-        inv_timescales = torch.reshape(inv_timescales, [batch_size, -1])
-        scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape(
-            inv_timescales, [1, 1, -1]
-        )
-        encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)
-        return encoding.type(dtype)
-
-    def forward(self, x):
-        batch_size, timesteps, input_dim = x.size()
-        positions = torch.arange(1, timesteps + 1, device=x.device)[None, :]
-        position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device)
-
-        return x + position_encoding
-
-
-class PositionwiseFeedForward(torch.nn.Module):
-    """Positionwise feed forward layer.
-
-    Args:
-        idim (int): Input dimenstion.
-        hidden_units (int): The number of hidden units.
-        dropout_rate (float): Dropout rate.
-
-    """
-
-    def __init__(self, idim, hidden_units, dropout_rate, activation=torch.nn.ReLU()):
-        """Construct an PositionwiseFeedForward object."""
-        super(PositionwiseFeedForward, self).__init__()
-        self.w_1 = torch.nn.Linear(idim, hidden_units)
-        self.w_2 = torch.nn.Linear(hidden_units, idim)
-        self.dropout = torch.nn.Dropout(dropout_rate)
-        self.activation = activation
-
-    def forward(self, x):
-        """Forward function."""
-        return self.w_2(self.dropout(self.activation(self.w_1(x))))
-
-
-class MultiHeadedAttentionSANM(nn.Module):
-    """Multi-Head Attention layer.
-
-    Args:
-        n_head (int): The number of heads.
-        n_feat (int): The number of features.
-        dropout_rate (float): Dropout rate.
-
-    """
-
-    def __init__(
-        self,
-        n_head,
-        in_feat,
-        n_feat,
-        dropout_rate,
-        kernel_size,
-        sanm_shfit=0,
-        lora_list=None,
-        lora_rank=8,
-        lora_alpha=16,
-        lora_dropout=0.1,
-    ):
-        """Construct an MultiHeadedAttention object."""
-        super().__init__()
-        assert n_feat % n_head == 0
-        # We assume d_v always equals d_k
-        self.d_k = n_feat // n_head
-        self.h = n_head
-        # self.linear_q = nn.Linear(n_feat, n_feat)
-        # self.linear_k = nn.Linear(n_feat, n_feat)
-        # self.linear_v = nn.Linear(n_feat, n_feat)
-
-        self.linear_out = nn.Linear(n_feat, n_feat)
-        self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3)
-        self.attn = None
-        self.dropout = nn.Dropout(p=dropout_rate)
-
-        self.fsmn_block = nn.Conv1d(
-            n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
-        )
-        # padding
-        left_padding = (kernel_size - 1) // 2
-        if sanm_shfit > 0:
-            left_padding = left_padding + sanm_shfit
-        right_padding = kernel_size - 1 - left_padding
-        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)
-
-    def forward_fsmn(self, inputs, mask, mask_shfit_chunk=None):
-        b, t, d = inputs.size()
-        if mask is not None:
-            mask = torch.reshape(mask, (b, -1, 1))
-            if mask_shfit_chunk is not None:
-                mask = mask * mask_shfit_chunk
-            inputs = inputs * mask
-
-        x = inputs.transpose(1, 2)
-        x = self.pad_fn(x)
-        x = self.fsmn_block(x)
-        x = x.transpose(1, 2)
-        x += inputs
-        x = self.dropout(x)
-        if mask is not None:
-            x = x * mask
-        return x
-
-    def forward_qkv(self, x):
-        """Transform query, key and value.
-
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-
-        Returns:
-            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
-            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
-            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
-
-        """
-        b, t, d = x.size()
-        q_k_v = self.linear_q_k_v(x)
-        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
-        q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(
-            1, 2
-        )  # (batch, head, time1, d_k)
-        k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(
-            1, 2
-        )  # (batch, head, time2, d_k)
-        v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(
-            1, 2
-        )  # (batch, head, time2, d_k)
-
-        return q_h, k_h, v_h, v
-
-    def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None):
-        """Compute attention context vector.
-
-        Args:
-            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
-            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
-            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
-
-        Returns:
-            torch.Tensor: Transformed value (#batch, time1, d_model)
-                weighted by the attention score (#batch, time1, time2).
-
-        """
-        n_batch = value.size(0)
-        if mask is not None:
-            if mask_att_chunk_encoder is not None:
-                mask = mask * mask_att_chunk_encoder
-
-            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
-
-            min_value = -float(
-                "inf"
-            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
-            scores = scores.masked_fill(mask, min_value)
-            attn = torch.softmax(scores, dim=-1).masked_fill(
-                mask, 0.0
-            )  # (batch, head, time1, time2)
-        else:
-            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
-
-        p_attn = self.dropout(attn)
-        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = (
-            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
-        )  # (batch, time1, d_model)
-
-        return self.linear_out(x)  # (batch, time1, d_model)
-
-    def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
-        """Compute scaled dot product attention.
-
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
-                (#batch, time1, time2).
-
-        Returns:
-            torch.Tensor: Output tensor (#batch, time1, d_model).
-
-        """
-        q_h, k_h, v_h, v = self.forward_qkv(x)
-        fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk)
-        q_h = q_h * self.d_k ** (-0.5)
-        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
-        att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
-        return att_outs + fsmn_memory
-
-    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
-        """Compute scaled dot product attention.
-
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
-                (#batch, time1, time2).
-
-        Returns:
-            torch.Tensor: Output tensor (#batch, time1, d_model).
-
-        """
-        q_h, k_h, v_h, v = self.forward_qkv(x)
-        if chunk_size is not None and look_back > 0 or look_back == -1:
-            if cache is not None:
-                k_h_stride = k_h[:, :, : -(chunk_size[2]), :]
-                v_h_stride = v_h[:, :, : -(chunk_size[2]), :]
-                k_h = torch.cat((cache["k"], k_h), dim=2)
-                v_h = torch.cat((cache["v"], v_h), dim=2)
-
-                cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2)
-                cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2)
-                if look_back != -1:
-                    cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :]
-                    cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :]
-            else:
-                cache_tmp = {
-                    "k": k_h[:, :, : -(chunk_size[2]), :],
-                    "v": v_h[:, :, : -(chunk_size[2]), :],
-                }
-                cache = cache_tmp
-        fsmn_memory = self.forward_fsmn(v, None)
-        q_h = q_h * self.d_k ** (-0.5)
-        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
-        att_outs = self.forward_attention(v_h, scores, None)
-        return att_outs + fsmn_memory, cache
-
-
-class LayerNorm(nn.LayerNorm):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def forward(self, input):
-        output = F.layer_norm(
-            input.float(),
-            self.normalized_shape,
-            self.weight.float() if self.weight is not None else None,
-            self.bias.float() if self.bias is not None else None,
-            self.eps,
-        )
-        return output.type_as(input)
-
-
-def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None):
-    if maxlen is None:
-        maxlen = lengths.max()
-    row_vector = torch.arange(0, maxlen, 1).to(lengths.device)
-    matrix = torch.unsqueeze(lengths, dim=-1)
-    mask = row_vector < matrix
-    mask = mask.detach()
-
-    return mask.to(dtype).to(device) if device is not None else mask.to(dtype)
-    # return mask.type(dtype).to(device) if device is not None else mask.type(dtype)
-
-
-class EncoderLayerSANM(nn.Module):
-    def __init__(
-        self,
-        in_size,
-        size,
-        self_attn,
-        feed_forward,
-        dropout_rate,
-        normalize_before=True,
-        concat_after=False,
-        stochastic_depth_rate=0.0,
-    ):
-        """Construct an EncoderLayer object."""
-        super(EncoderLayerSANM, self).__init__()
-        self.self_attn = self_attn
-        self.feed_forward = feed_forward
-        self.norm1 = LayerNorm(in_size)
-        self.norm2 = LayerNorm(size)
-        self.dropout = nn.Dropout(dropout_rate)
-        self.in_size = in_size
-        self.size = size
-        self.normalize_before = normalize_before
-        self.concat_after = concat_after
-        if self.concat_after:
-            self.concat_linear = nn.Linear(size + size, size)
-        self.stochastic_depth_rate = stochastic_depth_rate
-        self.dropout_rate = dropout_rate
-
-    def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
-        """Compute encoded features.
-
-        Args:
-            x_input (torch.Tensor): Input tensor (#batch, time, size).
-            mask (torch.Tensor): Mask tensor for the input (#batch, time).
-            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
-
-        Returns:
-            torch.Tensor: Output tensor (#batch, time, size).
-            torch.Tensor: Mask tensor (#batch, time).
-
-        """
-        skip_layer = False
-        # with stochastic depth, residual connection `x + f(x)` becomes
-        # `x <- x + 1 / (1 - p) * f(x)` at training time.
-        stoch_layer_coeff = 1.0
-        if self.training and self.stochastic_depth_rate > 0:
-            skip_layer = torch.rand(1).item() < self.stochastic_depth_rate
-            stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
-
-        if skip_layer:
-            if cache is not None:
-                x = torch.cat([cache, x], dim=1)
-            return x, mask
-
-        residual = x
-        if self.normalize_before:
-            x = self.norm1(x)
-
-        if self.concat_after:
-            x_concat = torch.cat(
-                (
-                    x,
-                    self.self_attn(
-                        x,
-                        mask,
-                        mask_shfit_chunk=mask_shfit_chunk,
-                        mask_att_chunk_encoder=mask_att_chunk_encoder,
-                    ),
-                ),
-                dim=-1,
-            )
-            if self.in_size == self.size:
-                x = residual + stoch_layer_coeff * self.concat_linear(x_concat)
-            else:
-                x = stoch_layer_coeff * self.concat_linear(x_concat)
-        else:
-            if self.in_size == self.size:
-                x = residual + stoch_layer_coeff * self.dropout(
-                    self.self_attn(
-                        x,
-                        mask,
-                        mask_shfit_chunk=mask_shfit_chunk,
-                        mask_att_chunk_encoder=mask_att_chunk_encoder,
-                    )
-                )
-            else:
-                x = stoch_layer_coeff * self.dropout(
-                    self.self_attn(
-                        x,
-                        mask,
-                        mask_shfit_chunk=mask_shfit_chunk,
-                        mask_att_chunk_encoder=mask_att_chunk_encoder,
-                    )
-                )
-        if not self.normalize_before:
-            x = self.norm1(x)
-
-        residual = x
-        if self.normalize_before:
-            x = self.norm2(x)
-        x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x))
-        if not self.normalize_before:
-            x = self.norm2(x)
-
-        return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder
-
-    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
-        """Compute encoded features.
-
-        Args:
-            x_input (torch.Tensor): Input tensor (#batch, time, size).
-            mask (torch.Tensor): Mask tensor for the input (#batch, time).
-            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
-
-        Returns:
-            torch.Tensor: Output tensor (#batch, time, size).
-            torch.Tensor: Mask tensor (#batch, time).
-
-        """
-
-        residual = x
-        if self.normalize_before:
-            x = self.norm1(x)
-
-        if self.in_size == self.size:
-            attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
-            x = residual + attn
-        else:
-            x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
-
-        if not self.normalize_before:
-            x = self.norm1(x)
-
-        residual = x
-        if self.normalize_before:
-            x = self.norm2(x)
-        x = residual + self.feed_forward(x)
-        if not self.normalize_before:
-            x = self.norm2(x)
-
-        return x, cache
-
-
-@tables.register("encoder_classes", "SenseVoiceEncoderSmall")
-class SenseVoiceEncoderSmall(nn.Module):
-    """
-    Author: Speech Lab of DAMO Academy, Alibaba Group
-    SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
-    https://arxiv.org/abs/2006.01713
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int = 256,
-        attention_heads: int = 4,
-        linear_units: int = 2048,
-        num_blocks: int = 6,
-        tp_blocks: int = 0,
-        dropout_rate: float = 0.1,
-        positional_dropout_rate: float = 0.1,
-        attention_dropout_rate: float = 0.0,
-        stochastic_depth_rate: float = 0.0,
-        input_layer: Optional[str] = "conv2d",
-        pos_enc_class=SinusoidalPositionEncoder,
-        normalize_before: bool = True,
-        concat_after: bool = False,
-        positionwise_layer_type: str = "linear",
-        positionwise_conv_kernel_size: int = 1,
-        padding_idx: int = -1,
-        kernel_size: int = 11,
-        sanm_shfit: int = 0,
-        selfattention_layer_type: str = "sanm",
-        **kwargs,
-    ):
-        super().__init__()
-        self._output_size = output_size
-
-        self.embed = SinusoidalPositionEncoder()
-
-        self.normalize_before = normalize_before
-
-        positionwise_layer = PositionwiseFeedForward
-        positionwise_layer_args = (
-            output_size,
-            linear_units,
-            dropout_rate,
-        )
-
-        encoder_selfattn_layer = MultiHeadedAttentionSANM
-        encoder_selfattn_layer_args0 = (
-            attention_heads,
-            input_size,
-            output_size,
-            attention_dropout_rate,
-            kernel_size,
-            sanm_shfit,
-        )
-        encoder_selfattn_layer_args = (
-            attention_heads,
-            output_size,
-            output_size,
-            attention_dropout_rate,
-            kernel_size,
-            sanm_shfit,
-        )
-
-        self.encoders0 = nn.ModuleList(
-            [
-                EncoderLayerSANM(
-                    input_size,
-                    output_size,
-                    encoder_selfattn_layer(*encoder_selfattn_layer_args0),
-                    positionwise_layer(*positionwise_layer_args),
-                    dropout_rate,
-                )
-                for i in range(1)
-            ]
-        )
-        self.encoders = nn.ModuleList(
-            [
-                EncoderLayerSANM(
-                    output_size,
-                    output_size,
-                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
-                    positionwise_layer(*positionwise_layer_args),
-                    dropout_rate,
-                )
-                for i in range(num_blocks - 1)
-            ]
-        )
-
-        self.tp_encoders = nn.ModuleList(
-            [
-                EncoderLayerSANM(
-                    output_size,
-                    output_size,
-                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
-                    positionwise_layer(*positionwise_layer_args),
-                    dropout_rate,
-                )
-                for i in range(tp_blocks)
-            ]
-        )
-
-        self.after_norm = LayerNorm(output_size)
-
-        self.tp_norm = LayerNorm(output_size)
-
-    def output_size(self) -> int:
-        return self._output_size
-
-    def forward(
-        self,
-        xs_pad: torch.Tensor,
-        ilens: torch.Tensor,
-    ):
-        """Embed positions in tensor."""
-        masks = sequence_mask(ilens, dtype=torch.bfloat16, device=ilens.device)[:, None, :]
-        # print(f"{masks=}")
-        # print(f"{ilens=}")
-        # print(f"{(masks>0.5).squeeze(1).sum(1).int()=}")
-
-        xs_pad *= self.output_size() ** 0.5
-
-        xs_pad = self.embed(xs_pad)
-
-        # forward encoder1
-        for layer_idx, encoder_layer in enumerate(self.encoders0):
-            encoder_outs = encoder_layer(xs_pad, masks)
-            xs_pad, masks = encoder_outs[0], encoder_outs[1]
-
-        for layer_idx, encoder_layer in enumerate(self.encoders):
-            encoder_outs = encoder_layer(xs_pad, masks)
-            xs_pad, masks = encoder_outs[0], encoder_outs[1]
-
-        xs_pad = self.after_norm(xs_pad)
-
-        # forward encoder2
-        # olens = masks.squeeze(1).sum(1).int()
-        olens = (masks > 0.5).squeeze(1).sum(1).int()
-
-        for layer_idx, encoder_layer in enumerate(self.tp_encoders):
-            encoder_outs = encoder_layer(xs_pad, masks)
-            xs_pad, masks = encoder_outs[0], encoder_outs[1]
-
-        xs_pad = self.tp_norm(xs_pad)
-        return xs_pad, olens
-
-
-@tables.register("model_classes", "SenseVoiceSmall")
-class SenseVoiceSmall(nn.Module):
-    """CTC-attention hybrid Encoder-Decoder model"""
-
-    def __init__(
-        self,
-        specaug: str = None,
-        specaug_conf: dict = None,
-        normalize: str = None,
-        normalize_conf: dict = None,
-        encoder: str = None,
-        encoder_conf: dict = None,
-        ctc_conf: dict = None,
-        input_size: int = 80,
-        vocab_size: int = -1,
-        ignore_id: int = -1,
-        blank_id: int = 0,
-        sos: int = 1,
-        eos: int = 2,
-        length_normalized_loss: bool = False,
-        **kwargs,
-    ):
-
-        super().__init__()
-
-        if specaug is not None:
-            specaug_class = tables.specaug_classes.get(specaug)
-            specaug = specaug_class(**specaug_conf)
-        if normalize is not None:
-            normalize_class = tables.normalize_classes.get(normalize)
-            normalize = normalize_class(**normalize_conf)
-        encoder_class = tables.encoder_classes.get(encoder)
-        encoder = encoder_class(input_size=input_size, **encoder_conf)
-        encoder_output_size = encoder.output_size()
-
-        if ctc_conf is None:
-            ctc_conf = {}
-        ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf)
-
-        self.blank_id = blank_id
-        self.sos = sos if sos is not None else vocab_size - 1
-        self.eos = eos if eos is not None else vocab_size - 1
-        self.vocab_size = vocab_size
-        self.ignore_id = ignore_id
-        self.specaug = specaug
-        self.normalize = normalize
-        self.encoder = encoder
-        self.error_calculator = None
-
-        self.ctc = ctc
-
-        self.length_normalized_loss = length_normalized_loss
-        self.encoder_output_size = encoder_output_size
-
-        self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13}
-        self.lid_int_dict = {24884: 3, 24885: 4, 24888: 7, 24892: 11, 24896: 12, 24992: 13}
-        self.textnorm_dict = {"withitn": 14, "woitn": 15}
-        self.textnorm_int_dict = {25016: 14, 25017: 15}
-        self.embed = torch.nn.Embedding(7 + len(self.lid_dict) + len(self.textnorm_dict), input_size)
-        self.emo_dict = {"unk": 25009, "happy": 25001, "sad": 25002, "angry": 25003, "neutral": 25004}
-        
-        self.criterion_att = LabelSmoothingLoss(
-            size=self.vocab_size,
-            padding_idx=self.ignore_id,
-            smoothing=kwargs.get("lsm_weight", 0.0),
-            normalize_length=self.length_normalized_loss,
-        )
-    
-    @staticmethod
-    def from_pretrained(model:str=None, **kwargs):
-        from funasr import AutoModel
-        model, kwargs = AutoModel.build_model(model=model, trust_remote_code=True, **kwargs)
-        
-        return model, kwargs
-
-    def forward(
-        self,
-        speech: torch.Tensor,
-        speech_lengths: torch.Tensor,
-        text: torch.Tensor,
-        text_lengths: torch.Tensor,
-        **kwargs,
-    ):
-        """Encoder + Decoder + Calc loss
-        Args:
-                speech: (Batch, Length, ...)
-                speech_lengths: (Batch, )
-                text: (Batch, Length)
-                text_lengths: (Batch,)
-        """
-        # import pdb;
-        # pdb.set_trace()
-        if len(text_lengths.size()) > 1:
-            text_lengths = text_lengths[:, 0]
-        if len(speech_lengths.size()) > 1:
-            speech_lengths = speech_lengths[:, 0]
-
-        batch_size = speech.shape[0]
-
-        # 1. Encoder
-        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, text)
-
-        loss_ctc, cer_ctc = None, None
-        loss_rich, acc_rich = None, None
-        stats = dict()
-
-        loss_ctc, cer_ctc = self._calc_ctc_loss(
-            encoder_out[:, 4:, :], encoder_out_lens - 4, text[:, 4:], text_lengths - 4
-        )
-
-        loss_rich, acc_rich = self._calc_rich_ce_loss(
-            encoder_out[:, :4, :], text[:, :4]
-        )
-
-        loss = loss_ctc + loss_rich
-        # Collect total loss stats
-        stats["loss_ctc"] = torch.clone(loss_ctc.detach()) if loss_ctc is not None else None
-        stats["loss_rich"] = torch.clone(loss_rich.detach()) if loss_rich is not None else None
-        stats["loss"] = torch.clone(loss.detach()) if loss is not None else None
-        stats["acc_rich"] = acc_rich
-
-        # force_gatherable: to-device and to-tensor if scalar for DataParallel
-        if self.length_normalized_loss:
-            batch_size = int((text_lengths + 1).sum())
-        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
-        return loss, stats, weight
-
-    def encode(
-        self,
-        speech: torch.Tensor,
-        speech_lengths: torch.Tensor,
-        text: torch.Tensor,
-        **kwargs,
-    ):
-        """Frontend + Encoder. Note that this method is used by asr_inference.py
-        Args:
-                speech: (Batch, Length, ...)
-                speech_lengths: (Batch, )
-                ind: int
-        """
-
-        # Data augmentation
-        if self.specaug is not None and self.training:
-            speech, speech_lengths = self.specaug(speech, speech_lengths)
-
-        # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
-        if self.normalize is not None:
-            speech, speech_lengths = self.normalize(speech, speech_lengths)
-
-
-        lids = torch.LongTensor([[self.lid_int_dict[int(lid)] 
-                                if torch.rand(1) > 0.2 and int(lid) in self.lid_int_dict else 0 ] 
-                                for lid in text[:, 0]]).to(speech.device)
-        language_query = self.embed(lids)
-        
-        styles = torch.LongTensor([[self.textnorm_int_dict[int(style)]] for style in text[:, 3]]).to(speech.device)
-        style_query = self.embed(styles)
-        speech = torch.cat((style_query, speech), dim=1)
-        speech_lengths += 1
-
-        event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat(speech.size(0), 1, 1)
-        input_query = torch.cat((language_query, event_emo_query), dim=1)
-        speech = torch.cat((input_query, speech), dim=1)
-        speech_lengths += 3
-
-        encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths)
-
-        return encoder_out, encoder_out_lens
-
-    def _calc_ctc_loss(
-        self,
-        encoder_out: torch.Tensor,
-        encoder_out_lens: torch.Tensor,
-        ys_pad: torch.Tensor,
-        ys_pad_lens: torch.Tensor,
-    ):
-        # Calc CTC loss
-        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)
-
-        # Calc CER using CTC
-        cer_ctc = None
-        if not self.training and self.error_calculator is not None:
-            ys_hat = self.ctc.argmax(encoder_out).data
-            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
-        return loss_ctc, cer_ctc
-
-    def _calc_rich_ce_loss(
-        self,
-        encoder_out: torch.Tensor,
-        ys_pad: torch.Tensor,
-    ):
-        decoder_out = self.ctc.ctc_lo(encoder_out)
-        # 2. Compute attention loss
-        loss_rich = self.criterion_att(decoder_out, ys_pad.contiguous())
-        acc_rich = th_accuracy(
-            decoder_out.view(-1, self.vocab_size),
-            ys_pad.contiguous(),
-            ignore_label=self.ignore_id,
-        )
-
-        return loss_rich, acc_rich
-
-
-    def inference(
-        self,
-        data_in,
-        data_lengths=None,
-        key: list = ["wav_file_tmp_name"],
-        tokenizer=None,
-        frontend=None,
-        **kwargs,
-    ):
-
-
-        meta_data = {}
-        if (
-            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
-        ):  # fbank
-            speech, speech_lengths = data_in, data_lengths
-            if len(speech.shape) < 3:
-                speech = speech[None, :, :]
-            if speech_lengths is None:
-                speech_lengths = speech.shape[1]
-        else:
-            # extract fbank feats
-            time1 = time.perf_counter()
-            audio_sample_list = load_audio_text_image_video(
-                data_in,
-                fs=frontend.fs,
-                audio_fs=kwargs.get("fs", 16000),
-                data_type=kwargs.get("data_type", "sound"),
-                tokenizer=tokenizer,
-            )
-            time2 = time.perf_counter()
-            meta_data["load_data"] = f"{time2 - time1:0.3f}"
-            speech, speech_lengths = extract_fbank(
-                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
-            )
-            time3 = time.perf_counter()
-            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
-            meta_data["batch_data_time"] = (
-                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
-            )
-
-        speech = speech.to(device=kwargs["device"])
-        speech_lengths = speech_lengths.to(device=kwargs["device"])
-
-        language = kwargs.get("language", "auto")
-        language_query = self.embed(
-            torch.LongTensor(
-                [[self.lid_dict[language] if language in self.lid_dict else 0]]
-            ).to(speech.device)
-        ).repeat(speech.size(0), 1, 1)
-        
-        use_itn = kwargs.get("use_itn", False)
-        output_timestamp = kwargs.get("output_timestamp", False)
-
-        textnorm = kwargs.get("text_norm", None)
-        if textnorm is None:
-            textnorm = "withitn" if use_itn else "woitn"
-        textnorm_query = self.embed(
-            torch.LongTensor([[self.textnorm_dict[textnorm]]]).to(speech.device)
-        ).repeat(speech.size(0), 1, 1)
-        speech = torch.cat((textnorm_query, speech), dim=1)
-        speech_lengths += 1
-
-        event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat(
-            speech.size(0), 1, 1
-        )
-        input_query = torch.cat((language_query, event_emo_query), dim=1)
-        speech = torch.cat((input_query, speech), dim=1)
-        speech_lengths += 3
-
-        # Encoder
-        encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths)
-        if isinstance(encoder_out, tuple):
-            encoder_out = encoder_out[0]
-
-        # c. Passed the encoder result and the beam search
-        ctc_logits = self.ctc.log_softmax(encoder_out)
-        if kwargs.get("ban_emo_unk", False):
-            ctc_logits[:, :, self.emo_dict["unk"]] = -float("inf")
-
-        results = []
-        b, n, d = encoder_out.size()
-        if isinstance(key[0], (list, tuple)):
-            key = key[0]
-        if len(key) < b:
-            key = key * b
-        for i in range(b):
-            x = ctc_logits[i, : encoder_out_lens[i].item(), :]
-            yseq = x.argmax(dim=-1)
-            yseq = torch.unique_consecutive(yseq, dim=-1)
-
-            ibest_writer = None
-            if kwargs.get("output_dir") is not None:
-                if not hasattr(self, "writer"):
-                    self.writer = DatadirWriter(kwargs.get("output_dir"))
-                ibest_writer = self.writer[f"1best_recog"]
-
-            mask = yseq != self.blank_id
-            token_int = yseq[mask].tolist()
-
-            # Change integer-ids to tokens
-            text = tokenizer.decode(token_int)
-            if ibest_writer is not None:
-                ibest_writer["text"][key[i]] = text
-
-            if output_timestamp:
-                from itertools import groupby
-                timestamp = []
-                tokens = tokenizer.text2tokens(text)[4:]
-
-                logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :]
-
-                pred = logits_speech.argmax(-1).cpu()
-                logits_speech[pred==self.blank_id, self.blank_id] = 0
-
-                align = ctc_forced_align(
-                    logits_speech.unsqueeze(0).float(),
-                    torch.Tensor(token_int[4:]).unsqueeze(0).long().to(logits_speech.device),
-                    (encoder_out_lens-4).long(),
-                    torch.tensor(len(token_int)-4).unsqueeze(0).long().to(logits_speech.device),
-                    ignore_id=self.ignore_id,
-                )
-
-                pred = groupby(align[0, :encoder_out_lens[0]])
-                _start = 0
-                token_id = 0
-                ts_max = encoder_out_lens[i] - 4
-                for pred_token, pred_frame in pred:
-                    _end = _start + len(list(pred_frame))
-                    if pred_token != 0:
-                        ts_left = max((_start*60-30)/1000, 0)
-                        ts_right = min((_end*60-30)/1000, (ts_max*60-30)/1000)
-                        timestamp.append([tokens[token_id], ts_left, ts_right])
-                        token_id += 1
-                    _start = _end
-
-                result_i = {"key": key[i], "text": text, "timestamp": timestamp}
-                results.append(result_i)
-            else:
-                result_i = {"key": key[i], "text": text}
-                results.append(result_i)
-        return results, meta_data
-
-
-    def inference_encode(
-        self,
-        data_in,
-        data_lengths=None,
-        key: list = ["wav_file_tmp_name"],
-        tokenizer=None,
-        frontend=None,
-        **kwargs,
-    ):
-
-        # fbank
-        speech, speech_lengths = data_in, data_lengths
-        if len(speech.shape) < 3:
-            speech = speech[None, :, :]
-        if speech_lengths is None:
-            speech_lengths = speech.shape[1]
-
-        speech = speech.to(device=kwargs["device"])
-        speech_lengths = speech_lengths.to(device=kwargs["device"])
-
-        language = kwargs.get("language", "auto")
-        language_query = self.embed(
-            torch.LongTensor(
-                [[self.lid_dict[language] if language in self.lid_dict else 0]]
-            ).to(speech.device)
-        ).repeat(speech.size(0), 1, 1)
-        
-        use_itn = kwargs.get("use_itn", False)
-        output_timestamp = kwargs.get("output_timestamp", False)
-
-        textnorm = kwargs.get("text_norm", None)
-        if textnorm is None:
-            textnorm = "withitn" if use_itn else "woitn"
-        textnorm_query = self.embed(
-            torch.LongTensor([[self.textnorm_dict[textnorm]]]).to(speech.device)
-        ).repeat(speech.size(0), 1, 1)
-        speech = torch.cat((textnorm_query, speech), dim=1)
-        speech_lengths += 1
-
-        event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat(
-            speech.size(0), 1, 1
-        )
-        input_query = torch.cat((language_query, event_emo_query), dim=1)
-        speech = torch.cat((input_query, speech), dim=1)
-        speech_lengths += 3
-
-        # Encoder
-        encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths)
-        if isinstance(encoder_out, tuple):
-            encoder_out = encoder_out[0]
-
-        return encoder_out, encoder_out_lens
-
-    def export(self, **kwargs):
-        from export_meta import export_rebuild_model
-
-        if "max_seq_len" not in kwargs:
-            kwargs["max_seq_len"] = 512
-        models = export_rebuild_model(model=self, **kwargs)
-        return models
-
-
-
-class AudioEncoder(nn.Module):
-
-    def __init__(
-        self,
-        model_dir=None
-    ):
-
-        super().__init__()
-
-        # model_dir = "SenceVoiceSmall"
-        from huggingface_hub import snapshot_download
-        model_dir = snapshot_download(repo_id="FunAudioLLM/SenseVoiceSmall")
-        self.model, self.kwargs = self.build_model(model=model_dir, trust_remote_code=True,)
-
-    
-    def forward(
-        self,
-        audios,
-    ):
-
-        # print(f"{[x for x in audios]=}")
-        # print(f"{[x.sum() for x in audios]=}")
-        # print(f"{[x.size() for x in audios]=}")
-
-        from torch.nn.utils.rnn import pad_sequence
-        feats_pad = pad_sequence(audios, batch_first=True, padding_value=0.0)
-        # feats_lens = torch.as_tensor([len(x) + 4 for x in audios])
-        feats_lens = torch.as_tensor([len(x) for x in audios])
-
-        feats_pad = feats_pad.to(torch.bfloat16)
-
-        # print(f"{feats_pad.sum()=}")
-        # print(f"{feats_pad.size()=}")
-        # print(f"{feats_lens=}")
-
-        encoder_out, encoder_out_lens = self.model.inference_encode(
-            feats_pad,
-            data_lengths=feats_lens,
-            language="auto", # "zh", "en", "yue", "ja", "ko", "nospeech"
-            use_itn=False,
-            ban_emo_unk=False,
-            **self.kwargs,
-        )
-
-        # print(f"{encoder_out.sum()=}")
-        # print(f"{encoder_out.size()=}")
-        # print(f"{encoder_out_lens=}")
-        return encoder_out, encoder_out_lens
-
-        audio_embeds = []
-        for x, y in zip(encoder_out, encoder_out_lens):
-            audio_embeds.append(x[:y, ...])
-
-        audio_embeds = torch.stack(audio_embeds, dim=0)
-
-        return audio_embeds
-
-    # https://github.com/modelscope/FunASR/blob/main/funasr/auto/auto_model.py
-    @staticmethod
-    def build_model(**kwargs):
-        from omegaconf import DictConfig, ListConfig
-        import os
-
-        from funasr.download.download_model_from_hub import download_model
-        from funasr.train_utils.set_all_random_seed import set_all_random_seed
-        from funasr.register import tables
-        from funasr.train_utils.load_pretrained_model import load_pretrained_model
-        from funasr.utils.misc import deep_update
-
-        import logging
-
-        assert "model" in kwargs
-        if "model_conf" not in kwargs:
-            logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
-            kwargs = download_model(**kwargs)
-
-        set_all_random_seed(kwargs.get("seed", 0))
-
-        device = kwargs.get("device", "cuda")
-        if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
-            device = "cpu"
-            kwargs["batch_size"] = 1
-        kwargs["device"] = device
-
-        torch.set_num_threads(kwargs.get("ncpu", 4))
-
-        # build tokenizer
-        tokenizer = kwargs.get("tokenizer", None)
-        kwargs["tokenizer"] = tokenizer
-        kwargs["vocab_size"] = -1
-
-        if tokenizer is not None:
-            tokenizers = (
-                tokenizer.split(",") if isinstance(tokenizer, str) else tokenizer
-            )  # type of tokenizers is list!!!
-            tokenizers_conf = kwargs.get("tokenizer_conf", {})
-            tokenizers_build = []
-            vocab_sizes = []
-            token_lists = []
-
-            ### === only for kws ===
-            token_list_files = kwargs.get("token_lists", [])
-            seg_dicts = kwargs.get("seg_dicts", [])
-            ### === only for kws ===
-
-            if not isinstance(tokenizers_conf, (list, tuple, ListConfig)):
-                tokenizers_conf = [tokenizers_conf] * len(tokenizers)
-
-            for i, tokenizer in enumerate(tokenizers):
-                tokenizer_class = tables.tokenizer_classes.get(tokenizer)
-                tokenizer_conf = tokenizers_conf[i]
-
-                ### === only for kws ===
-                if len(token_list_files) > 1:
-                    tokenizer_conf["token_list"] = token_list_files[i]
-                if len(seg_dicts) > 1:
-                    tokenizer_conf["seg_dict"] = seg_dicts[i]
-                ### === only for kws ===
-
-                tokenizer = tokenizer_class(**tokenizer_conf)
-                tokenizers_build.append(tokenizer)
-                token_list = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
-                token_list = (
-                    tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else token_list
-                )
-                vocab_size = -1
-                if token_list is not None:
-                    vocab_size = len(token_list)
-
-                if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"):
-                    vocab_size = tokenizer.get_vocab_size()
-                token_lists.append(token_list)
-                vocab_sizes.append(vocab_size)
-
-            if len(tokenizers_build) <= 1:
-                tokenizers_build = tokenizers_build[0]
-                token_lists = token_lists[0]
-                vocab_sizes = vocab_sizes[0]
-
-            kwargs["tokenizer"] = tokenizers_build
-            kwargs["vocab_size"] = vocab_sizes
-            kwargs["token_list"] = token_lists
-
-        # build frontend
-        frontend = kwargs.get("frontend", None)
-        kwargs["input_size"] = None
-        if frontend is not None:
-            frontend_class = tables.frontend_classes.get(frontend)
-            frontend = frontend_class(**kwargs.get("frontend_conf", {}))
-            kwargs["input_size"] = (
-                frontend.output_size() if hasattr(frontend, "output_size") else None
-            )
-        kwargs["frontend"] = frontend
-        # build model
-        model_class = tables.model_classes.get(kwargs["model"])
-        assert model_class is not None, f'{kwargs["model"]} is not registered'
-        model_conf = {}
-        deep_update(model_conf, kwargs.get("model_conf", {}))
-        deep_update(model_conf, kwargs)
-        model = model_class(**model_conf)
-
-        # init_param
-        init_param = kwargs.get("init_param", None)
-        if init_param is not None:
-            if os.path.exists(init_param):
-                logging.info(f"Loading pretrained params from {init_param}")
-                load_pretrained_model(
-                    model=model,
-                    path=init_param,
-                    ignore_init_mismatch=kwargs.get("ignore_init_mismatch", True),
-                    oss_bucket=kwargs.get("oss_bucket", None),
-                    scope_map=kwargs.get("scope_map", []),
-                    excludes=kwargs.get("excludes", None),
-                )
-            else:
-                print(f"error, init_param does not exist!: {init_param}")
-
-        # fp16
-        if kwargs.get("fp16", False):
-            model.to(torch.float16)
-        elif kwargs.get("bf16", False):
-            model.to(torch.bfloat16)
-        # model.to(device)
-
-        if not kwargs.get("disable_log", True):
-            tables.print()
-
-        return model, kwargs
-

omni_diffusion/models/dream/resampler_projector.py

@@ -1,40 +0,0 @@
-import torch
-import torch.nn as nn
-
-import math
-
-class ResamplerProjector(nn.Module):
-    """
-    A projection module designed to map input features (e.g., from a vision encoder)
-    to a target hidden size. It flattens spatial dimensions into a sequence 
-    and applies an MLP.
-    """
-    def __init__(self, proj_input_size, hidden_size):
-        super().__init__()
-
-        self.pre_proj_layernorm = torch.nn.LayerNorm(proj_input_size)
-
-        self.mlp = nn.Sequential(
-            nn.Linear(proj_input_size, hidden_size, bias=False),
-            nn.GELU(),
-            nn.Linear(hidden_size, hidden_size, bias=False),
-        )
-        self.mlp.apply(init_weights)
-        self.pre_proj_layernorm.apply(init_weights)
-
-    def forward(self, x, *args, **kwargs):
-        x = x.reshape(x.shape[0], -1, x.shape[-1])
-        x = self.pre_proj_layernorm(x)
-        x = self.mlp(x)
-        return x
-
-def init_weights(m):
-    if isinstance(m, nn.Linear):
-        torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
-        if m.bias is not None:
-            torch.nn.init.zeros_(m.bias)
-
-    if isinstance(m, nn.LayerNorm):
-        torch.nn.init.ones_(m.weight)
-        torch.nn.init.zeros_(m.bias)
-

omni_diffusion/models/dream/tokenization_dream.py

@@ -1,351 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and The HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on Qwen's implementations in this library.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for Dream."""
-
-import json
-import os
-import unicodedata
-from functools import lru_cache
-from typing import Optional, Tuple
-
-import regex as re
-
-from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
-    "vocab_file": "vocab.json",
-    "merges_file": "merges.txt",
-}
-
-
-MAX_MODEL_INPUT_SIZES = {"dream/dream-tokenizer": 32768}
-
-regex_parts = [
-    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)",       # Contractions (case-insensitive)
-    r"[^\r\n\p{L}\p{N}]?\p{L}+",           # Words (with optional leading symbol)
-    r"\p{N}",                              # Numbers
-    r" ?[^\s\p{L}\p{N}]+[\r\n]*",          # Punctuation and symbols
-    r"\s*[\r\n]+",                         # Newlines (handling indentation)
-    r"\s+(?!\S)",                          # Trailing whitespace
-    r"\s+",                                # General whitespace
-]
-PRETOKENIZE_REGEX = "|".join(regex_parts)
-
-
-@lru_cache()
-# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
-def bytes_to_unicode():
-    """
-    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
-    characters the bpe code barfs on.
-
-    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
-    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
-    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
-    tables between utf-8 bytes and unicode strings.
-    """
-    bs = (
-        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
-    )
-    cs = bs[:]
-    n = 0
-    for b in range(2**8):
-        if b not in bs:
-            bs.append(b)
-            cs.append(2**8 + n)
-            n += 1
-    cs = [chr(n) for n in cs]
-    return dict(zip(bs, cs))
-
-
-# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
-def get_pairs(word):
-    """
-    Return set of symbol pairs in a word.
-
-    Word is represented as tuple of symbols (symbols being variable-length strings).
-    """
-    pairs = set()
-    prev_char = word[0]
-    for char in word[1:]:
-        pairs.add((prev_char, char))
-        prev_char = char
-    return pairs
-
-
-class DreamTokenizer(PreTrainedTokenizer):
-    """
-    Construct a Dream tokenizer. Based on byte-level Byte-Pair-Encoding.
-
-    Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
-    be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
-    ```python
-    >>> from transformers import AutoTokenizer
-
-    >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-v0-Base-7B", trust_remote_code=True)
-    >>> tokenizer("Hello world")["input_ids"]
-    [9707, 1879]
-
-    >>> tokenizer(" Hello world")["input_ids"]
-    [21927, 1879]
-    ```
-    This is expected.
-
-    You should not use GPT2Tokenizer instead, because of the different pretokenization rules.
-
-    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
-    this superclass for more information regarding those methods.
-
-    Args:
-        vocab_file (`str`):
-            Path to the vocabulary file.
-        merges_file (`str`):
-            Path to the merges file.
-        errors (`str`, *optional*, defaults to `"replace"`):
-            Paradigm to follow when decoding bytes to UTF-8. See
-            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
-        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
-            token instead.
-        bos_token (`str`, *optional*):
-            The beginning of sequence token. Not applicable for this tokenizer.
-        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The end of sequence token.
-        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The token used for padding, for example when batching sequences of different lengths.
-        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
-            Whether or not the model should cleanup the spaces that were added when splitting the input text during the
-            tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.
-        split_special_tokens (`bool`, *optional*, defaults to `False`):
-            Whether or not the special tokens should be split during the tokenization process. The default behavior is
-            to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then 
-            `tokenizer.tokenize("<|endoftext|>") = ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, 
-            then `tokenizer.tokenize("<|endoftext|>")` will be give `['<', '|', 'endo', 'ft', 'ext', '|', '>']`. 
-            This argument is only supported for `slow` tokenizers for the moment.
-    """
-
-    vocab_files_names = VOCAB_FILES_NAMES
-    model_input_names = ["input_ids", "attention_mask"]
-
-    def __init__(
-        self,
-        vocab_file,
-        merges_file,
-        errors="replace",
-        unk_token="<|endoftext|>",
-        bos_token=None,
-        eos_token="<|endoftext|>",
-        pad_token="<|endoftext|>",
-        clean_up_tokenization_spaces=False,
-        split_special_tokens=False,
-        **kwargs,
-    ):
-        # Dream vocab does not contain control tokens; added tokens need to be special
-        bos_token = (
-            AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(bos_token, str)
-            else bos_token
-        )
-        eos_token = (
-            AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(eos_token, str)
-            else eos_token
-        )
-        unk_token = (
-            AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(unk_token, str)
-            else unk_token
-        )
-        pad_token = (
-            AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(pad_token, str)
-            else pad_token
-        )
-
-        with open(vocab_file, encoding="utf-8") as vocab_handle:
-            self.encoder = json.load(vocab_handle)
-        self.decoder = {v: k for k, v in self.encoder.items()}
-        self.errors = errors  # how to handle errors in decoding
-        self.byte_encoder = bytes_to_unicode()
-        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
-        bpe_merges = []
-        with open(merges_file, encoding="utf-8") as merges_handle:
-            for i, line in enumerate(merges_handle):
-                line = line.strip()
-                if (i == 0 and line.startswith("#version:")) or not line:
-                    continue
-                bpe_merges.append(tuple(line.split()))
-        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
-        # NOTE: the cache can grow without bound and will get really large for long running processes
-        # (esp. for texts of language that do not use space between word, e.g. Chinese); technically
-        # not a memory leak but appears as one.
-        # GPT2Tokenizer has the same problem, so let's be consistent.
-        self.cache = {}
-
-        self.pat = re.compile(PRETOKENIZE_REGEX)
-
-        if kwargs.get("add_prefix_space", False):
-            logger.warning_once(
-                f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."
-            )
-
-        super().__init__(
-            errors=errors,
-            bos_token=bos_token,
-            eos_token=eos_token,
-            pad_token=pad_token,
-            unk_token=unk_token,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-            split_special_tokens=split_special_tokens,
-            **kwargs,
-        )
-
-    @property
-    def vocab_size(self) -> int:
-        return len(self.encoder)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
-    def get_vocab(self):
-        return dict(self.encoder, **self.added_tokens_encoder)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
-    def bpe(self, token):
-        if token in self.cache:
-            return self.cache[token]
-        word = tuple(token)
-        pairs = get_pairs(word)
-
-        if not pairs:
-            return token
-
-        while True:
-            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
-            if bigram not in self.bpe_ranks:
-                break
-            first, second = bigram
-            new_word = []
-            i = 0
-            while i < len(word):
-                try:
-                    j = word.index(first, i)
-                except ValueError:
-                    new_word.extend(word[i:])
-                    break
-                else:
-                    new_word.extend(word[i:j])
-                    i = j
-
-                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
-                    new_word.append(first + second)
-                    i += 2
-                else:
-                    new_word.append(word[i])
-                    i += 1
-            new_word = tuple(new_word)
-            word = new_word
-            if len(word) == 1:
-                break
-            else:
-                pairs = get_pairs(word)
-        word = " ".join(word)
-        self.cache[token] = word
-        return word
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
-    def _tokenize(self, text):
-        """Tokenize a string."""
-        bpe_tokens = []
-        for token in re.findall(self.pat, text):
-            token = "".join(
-                self.byte_encoder[b] for b in token.encode("utf-8")
-            )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
-            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
-        return bpe_tokens
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
-    def _convert_token_to_id(self, token):
-        """Converts a token (str) in an id using the vocab."""
-        return self.encoder.get(token, self.encoder.get(self.unk_token))
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
-    def _convert_id_to_token(self, index):
-        """Converts an index (integer) in a token (str) using the vocab."""
-        return self.decoder.get(index)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
-    def convert_tokens_to_string(self, tokens):
-        """Converts a sequence of tokens (string) in a single string."""
-        text = "".join(tokens)
-        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
-        return text
-
-    def decode(
-        self,
-        token_ids,
-        skip_special_tokens: bool = False,
-        clean_up_tokenization_spaces: Optional[bool] = False,
-        spaces_between_special_tokens: bool = False,
-        **kwargs,
-    ) -> str:
-        # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers
-        # and cannot be configured elsewhere, but it should default to False for DreamTokenizer
-        return super().decode(
-            token_ids,
-            skip_special_tokens=skip_special_tokens,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-            spaces_between_special_tokens=spaces_between_special_tokens,
-            **kwargs,
-        )
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
-    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
-        if not os.path.isdir(save_directory):
-            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
-            return
-        vocab_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
-        )
-        merge_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
-        )
-
-        with open(vocab_file, "w", encoding="utf-8") as f:
-            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
-        index = 0
-        with open(merge_file, "w", encoding="utf-8") as writer:
-            writer.write("#version: 0.2\n")
-            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
-                if index != token_index:
-                    logger.warning(
-                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
-                        " Please check that the tokenizer is not corrupted!"
-                    )
-                    index = token_index
-                writer.write(" ".join(bpe_tokens) + "\n")
-                index += 1
-
-        return vocab_file, merge_file
-
-    def prepare_for_tokenization(self, text, **kwargs):
-        text = unicodedata.normalize("NFC", text)
-        return (text, kwargs)
-    

omni_diffusion/models/dream/tokenizer_config.json

@@ -1,219 +0,0 @@
-{
-  "add_bos_token": false,
-  "add_prefix_space": false,
-  "added_tokens_decoder": {
-    "151643": {
-      "content": "<|endoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151644": {
-      "content": "<|im_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151645": {
-      "content": "<|im_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151646": {
-      "content": "<|object_ref_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151647": {
-      "content": "<|object_ref_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151648": {
-      "content": "<|box_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151649": {
-      "content": "<|box_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151650": {
-      "content": "<|quad_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151651": {
-      "content": "<|quad_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151652": {
-      "content": "<|vision_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151653": {
-      "content": "<|vision_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151654": {
-      "content": "<|vision_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151655": {
-      "content": "<|image_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151656": {
-      "content": "<|video_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151657": {
-      "content": "<tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151658": {
-      "content": "</tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151659": {
-      "content": "<|fim_prefix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151660": {
-      "content": "<|fim_middle|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151661": {
-      "content": "<|fim_suffix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151662": {
-      "content": "<|fim_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151663": {
-      "content": "<|repo_name|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151664": {
-      "content": "<|file_sep|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151665": {
-      "content": "<|beginoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151666": {
-      "content": "<|mask|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "additional_special_tokens": [
-    "<|beginoftext|>",
-    "<|mask|>"
-  ],
-  "auto_map": {
-    "AutoTokenizer": [
-      "tokenization_dream.DreamTokenizer",
-      null
-    ]
-  },
-  "bos_token": "<|beginoftext|>",
-  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|endoftext|>",
-  "errors": "replace",
-  "mask_token": "<|mask|>",
-  "model_max_length": 131072,
-  "pad_token": "<|endoftext|>",
-  "split_special_tokens": false,
-  "tokenizer_class": "DreamTokenizer",
-  "unk_token": null
-}

omni_diffusion/models/magvit/common_modules.py

@@ -1,393 +0,0 @@
-"""
-Modified from https://github.com/CompVis/taming-transformers/blob/master/taming/modules/diffusionmodules/model.py#L34
-"""
-
-import math
-from typing import Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange, repeat
-from einops.layers.torch import Rearrange
-
-
-def nonlinearity(x):
-    # swish
-    return x * torch.sigmoid(x)
-
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(
-        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-    )
-
-
-class Upsample(nn.Module):
-    """
-    Spatial upsampling layer.
-    Upsamples by a factor of 2 using nearest neighbor interpolation,
-    optionally followed by a convolution.
-    """
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(
-                in_channels, in_channels, kernel_size=3, stride=1, padding=1
-            )
-
-    def forward(self, x):
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-
-class DepthToSpaceUpsample(nn.Module):
-    """
-    Upsamples spatially by moving channel depth to spatial dimensions (PixelShuffle).
-    Expands channels by 4, then rearranges (c * 2 * 2) -> (h * 2, w * 2).
-    """
-    def __init__(
-        self,
-        in_channels,
-    ):
-        super().__init__()
-        conv = nn.Conv2d(in_channels, in_channels * 4, 1)
-
-        self.net = nn.Sequential(
-            conv,
-            nn.SiLU(),
-            Rearrange("b (c p1 p2) h w -> b c (h p1) (w p2)", p1=2, p2=2),
-        )
-
-        self.init_conv_(conv)
-
-    def init_conv_(self, conv):
-        o, i, h, w = conv.weight.shape
-        conv_weight = torch.empty(o // 4, i, h, w)
-        nn.init.kaiming_uniform_(conv_weight)
-        conv_weight = repeat(conv_weight, "o ... -> (o 4) ...")
-
-        conv.weight.data.copy_(conv_weight)
-        nn.init.zeros_(conv.bias.data)
-
-    def forward(self, x):
-        out = self.net(x)
-        return out
-
-
-class Downsample(nn.Module):
-    """
-    Spatial downsampling layer.
-    Uses either a strided convolution (with padding) or average pooling.
-    """
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(
-                in_channels, in_channels, kernel_size=3, stride=2, padding=0
-            )
-
-    def forward(self, x):
-        if self.with_conv:
-            pad = (0, 1, 0, 1)
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-
-def unpack_time(t, batch):
-    _, c, w, h = t.size()
-    out = torch.reshape(t, [batch, -1, c, w, h])
-    out = rearrange(out, "b t c h w -> b c t h w")
-    return out
-
-
-def pack_time(t):
-    out = rearrange(t, "b c t h w -> b t c h w")
-    _, _, c, w, h = out.size()
-    return torch.reshape(out, [-1, c, w, h])
-
-
-class TimeDownsample2x(nn.Module):
-    """
-    Temporal downsampling layer.
-    Reduces the time dimension (T) by half using 1D convolution with causal padding.
-    """
-    def __init__(
-        self,
-        dim,
-        dim_out=None,
-        kernel_size=3,
-    ):
-        super().__init__()
-        if dim_out is None:
-            dim_out = dim
-        self.time_causal_padding = (kernel_size - 1, 0)
-        self.conv = nn.Conv1d(dim, dim_out, kernel_size, stride=2)
-
-    def forward(self, x):
-        x = rearrange(x, "b c t h w -> b h w c t")
-        b, h, w, c, t = x.size()
-        x = torch.reshape(x, [-1, c, t])
-
-        x = F.pad(x, self.time_causal_padding)
-        out = self.conv(x)
-
-        out = torch.reshape(out, [b, h, w, c, t])
-        out = rearrange(out, "b h w c t -> b c t h w")
-        out = rearrange(out, "b h w c t -> b c t h w")
-        return out
-
-
-class TimeUpsample2x(nn.Module):
-    """
-    Temporal upsampling layer.
-    Doubles the time dimension (T) using a PixelShuffle-like approach on the time axis.
-    """
-    def __init__(self, dim, dim_out=None):
-        super().__init__()
-        if dim_out is None:
-            dim_out = dim
-        conv = nn.Conv1d(dim, dim_out * 2, 1)
-
-        self.net = nn.Sequential(
-            nn.SiLU(), conv, Rearrange("b (c p) t -> b c (t p)", p=2)
-        )
-
-        self.init_conv_(conv)
-
-    def init_conv_(self, conv):
-        o, i, t = conv.weight.shape
-        conv_weight = torch.empty(o // 2, i, t)
-        nn.init.kaiming_uniform_(conv_weight)
-        conv_weight = repeat(conv_weight, "o ... -> (o 2) ...")
-
-        conv.weight.data.copy_(conv_weight)
-        nn.init.zeros_(conv.bias.data)
-
-    def forward(self, x):
-        x = rearrange(x, "b c t h w -> b h w c t")
-        b, h, w, c, t = x.size()
-        x = torch.reshape(x, [-1, c, t])
-
-        out = self.net(x)
-        out = out[:, :, 1:].contiguous()
-
-        out = torch.reshape(out, [b, h, w, c, t])
-        out = rearrange(out, "b h w c t -> b c t h w")
-        return out
-
-
-class AttnBlock(nn.Module):
-    """
-    Standard Self-Attention Block for spatial dimensions.
-    """
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.k = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.v = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.proj_out = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, h, w = q.shape
-        q = q.reshape(b, c, h * w)
-        q = q.permute(0, 2, 1)  # b,hw,c
-        k = k.reshape(b, c, h * w)  # b,c,hw
-        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b, c, h * w)
-        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b, c, h, w)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class TimeAttention(AttnBlock):
-    """
-    Applies the AttnBlock over the temporal dimension (T).
-    It treats spatial dimensions (H, W) as part of the batch.
-    """
-    def forward(self, x, *args, **kwargs):
-        x = rearrange(x, "b c t h w -> b h w t c")
-        b, h, w, t, c = x.size()
-        x = torch.reshape(x, (-1, t, c))
-
-        x = super().forward(x, *args, **kwargs)
-
-        x = torch.reshape(x, [b, h, w, t, c])
-        return rearrange(x, "b h w t c -> b c t h w")
-
-
-class Residual(nn.Module):
-    def __init__(self, fn: nn.Module):
-        super().__init__()
-        self.fn = fn
-
-    def forward(self, x, **kwargs):
-        return self.fn(x, **kwargs) + x
-
-
-def cast_tuple(t, length=1):
-    return t if isinstance(t, tuple) else ((t,) * length)
-
-
-class CausalConv3d(nn.Module):
-    """
-    3D Convolution with Causal Padding on the time axis.
-    Ensures that the convolution at time T only depends on frames <= T.
-    """
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int, int, int]],
-        pad_mode="constant",
-        **kwargs
-    ):
-        super().__init__()
-        kernel_size = cast_tuple(kernel_size, 3)
-
-        time_kernel_size, height_kernel_size, width_kernel_size = kernel_size
-
-        dilation = kwargs.pop("dilation", 1)
-        stride = kwargs.pop("stride", 1)
-
-        self.pad_mode = pad_mode
-        time_pad = dilation * (time_kernel_size - 1) + (1 - stride)
-        height_pad = height_kernel_size // 2
-        width_pad = width_kernel_size // 2
-
-        self.time_pad = time_pad
-        self.time_causal_padding = (
-            width_pad,
-            width_pad,
-            height_pad,
-            height_pad,
-            time_pad,
-            0,
-        )
-
-        stride = (stride, 1, 1)
-        dilation = (dilation, 1, 1)
-        self.conv = nn.Conv3d(
-            chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs
-        )
-
-    def forward(self, x):
-        pad_mode = self.pad_mode if self.time_pad < x.shape[2] else "constant"
-
-        x = F.pad(x, self.time_causal_padding, mode=pad_mode)
-        return self.conv(x)
-
-
-def ResnetBlockCausal3D(
-    dim, kernel_size: Union[int, Tuple[int, int, int]], pad_mode: str = "constant"
-):
-    net = nn.Sequential(
-        Normalize(dim),
-        nn.SiLU(),
-        CausalConv3d(dim, dim, kernel_size, pad_mode),
-        Normalize(dim),
-        nn.SiLU(),
-        CausalConv3d(dim, dim, kernel_size, pad_mode),
-    )
-    return Residual(net)
-
-
-class ResnetBlock(nn.Module):
-    """
-    Standard ResNet Block commonly used in Diffusion Models.
-    Includes support for Time Embeddings (temb) injection.
-    """
-    def __init__(
-        self,
-        *,
-        in_channels,
-        out_channels=None,
-        conv_shortcut=False,
-        dropout,
-        temb_channels=512
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(
-            in_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
-        else:
-            self.temb_proj = None
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(
-            out_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
-                )
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
-                )
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
-
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x + h

omni_diffusion/models/magvit/misc.py

@@ -1,55 +0,0 @@
-# copy from MMaDA
-
-from omegaconf import OmegaConf
-import torch
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    NewType,
-    Optional,
-    Sized,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-)
-try:
-    from typing import Literal
-except ImportError:
-    from typing_extensions import Literal
-
-# Tensor dtype
-# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
-from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
-
-# Config type
-from omegaconf import DictConfig
-
-# PyTorch Tensor type
-from torch import Tensor
-
-# Runtime type checking decorator
-from typeguard import typechecked as typechecker
-
-
-def broadcast(tensor, src=0):
-    if not _distributed_available():
-        return tensor
-    else:
-        torch.distributed.broadcast(tensor, src=src)
-        return tensor
-
-def _distributed_available():
-    return torch.distributed.is_available() and torch.distributed.is_initialized()
-
-def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
-    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
-    if '--local-rank' in cfg:
-        del cfg['--local-rank']
-    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
-    scfg = OmegaConf.structured(fields(**cfg))
-    return scfg

omni_diffusion/models/magvit/modeling_magvitv2.py

@@ -1,445 +0,0 @@
-"""
-Modified from https://github.com/Gen-Verse/MMaDA/blob/main/models/modeling_magvitv2.py
-"""
-
-from dataclasses import dataclass, field
-import numpy as np
-import torch
-import torch.nn as nn
-from .common_modules import *
-from .modeling_utils import ConfigMixin, ModelMixin, register_to_config
-from .misc import *
-import math
-
-class Updateable:
-    def do_update_step(
-            self, epoch: int, global_step: int, on_load_weights: bool = False
-    ):
-        for attr in self.__dir__():
-            if attr.startswith("_"):
-                continue
-            try:
-                module = getattr(self, attr)
-            except:
-                continue  # ignore attributes like property, which can't be retrived using getattr?
-            if isinstance(module, Updateable):
-                module.do_update_step(
-                    epoch, global_step, on_load_weights=on_load_weights
-                )
-        self.update_step(epoch, global_step, on_load_weights=on_load_weights)
-
-    def do_update_step_end(self, epoch: int, global_step: int):
-        for attr in self.__dir__():
-            if attr.startswith("_"):
-                continue
-            try:
-                module = getattr(self, attr)
-            except:
-                continue  # ignore attributes like property, which can't be retrived using getattr?
-            if isinstance(module, Updateable):
-                module.do_update_step_end(epoch, global_step)
-        self.update_step_end(epoch, global_step)
-
-    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
-        # override this method to implement custom update logic
-        # if on_load_weights is True, you should be careful doing things related to model evaluations,
-        # as the models and tensors are not guarenteed to be on the same device
-        pass
-
-    def update_step_end(self, epoch: int, global_step: int):
-        pass
-
-class VQGANEncoder(ModelMixin, ConfigMixin):
-    @dataclass
-    class Config:
-        ch: int = 128
-        ch_mult: List[int] = field(default_factory=lambda: [1, 2, 2, 4, 4])
-        num_res_blocks: List[int] = field(default_factory=lambda: [4, 3, 4, 3, 4])
-        attn_resolutions: List[int] = field(default_factory=lambda: [5])
-        dropout: float = 0.0
-        in_ch: int = 3
-        out_ch: int = 3
-        resolution: int = 256
-        z_channels: int = 13
-        double_z: bool = False
-
-    def __init__(self,
-                 ch: int = 128,
-                 ch_mult: List[int] = [1, 2, 2, 4, 4],
-                 num_res_blocks: List[int] = [4, 3, 4, 3, 4],
-                 attn_resolutions: List[int] = [5],
-                 dropout: float = 0.0,
-                 in_ch: int = 3,
-                 out_ch: int = 3,
-                 resolution: int = 256,
-                 z_channels: int = 13,
-                 double_z: bool = False):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_ch = in_ch
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(
-            self.in_ch, self.ch, kernel_size=3, stride=1, padding=1
-        )
-
-        curr_res = self.resolution
-        in_ch_mult = (1,) + tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = self.ch * in_ch_mult[i_level]
-            block_out = self.ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks[i_level]):
-                block.append(
-                    ResnetBlock(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        temb_channels=self.temb_ch,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                down.downsample = Downsample(block_in, True)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-
-
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(
-            block_in,
-            2 * z_channels if double_z else z_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-        )
-
-        self.quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
-        # for param in self.parameters():
-        #     broadcast(param, src=0)
-
-    def forward(self, x):
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks[i_level]):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions - 1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        h = self.quant_conv(h)
-        return h
-
-
-class LFQuantizer(nn.Module):
-    def __init__(self, num_codebook_entry: int = -1,
-                 codebook_dim: int = 13,
-                 beta: float = 0.25,
-                 entropy_multiplier: float = 0.1,
-                 commit_loss_multiplier: float = 0.1, ):
-        super().__init__()
-        self.codebook_size = 2 ** codebook_dim
-        print(
-            f"Look-up free quantizer with codebook size: {self.codebook_size}"
-        )
-        self.e_dim = codebook_dim
-        self.beta = beta
-
-        indices = torch.arange(self.codebook_size)
-
-        binary = (
-                         indices.unsqueeze(1)
-                         >> torch.arange(codebook_dim - 1, -1, -1, dtype=torch.long)
-                 ) & 1
-
-        embedding = binary.float() * 2 - 1
-        self.register_buffer("embedding", embedding)
-        self.register_buffer(
-            "power_vals", 2 ** torch.arange(codebook_dim - 1, -1, -1)
-        )
-        self.commit_loss_multiplier = commit_loss_multiplier
-        self.entropy_multiplier = entropy_multiplier
-
-    def get_indices(self, z_q):
-        return (
-            (self.power_vals.reshape(1, -1, 1, 1) * (z_q > 0).float())
-            .sum(1, keepdim=True)
-            .long()
-        )
-
-    def get_codebook_entry(self, indices, shape=None):
-        if shape is None:
-            h, w = int(math.sqrt(indices.shape[-1])), int(math.sqrt(indices.shape[-1]))
-        else:
-            h, w = shape
-        b, _ = indices.shape
-        indices = indices.reshape(-1)
-        z_q = self.embedding[indices]
-        z_q = z_q.view(b, h, w, -1)
-
-        # reshape back to match original input shape
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q
-
-    def forward(self, z, get_code=False):
-        """
-        Inputs the output of the encoder network z and maps it to a discrete
-        one-hot vector that is the index of the closest embedding vector e_j
-        z (continuous) -> z_q (discrete)
-        z.shape = (batch, channel, height, width)
-        quantization pipeline:
-            1. get encoder input (B,C,H,W)
-            2. flatten input to (B*H*W,C)
-        """
-        if get_code:
-            return self.get_codebook_entry(z)
-
-        # reshape z -> (batch, height, width, channel) and flatten
-        z = z.permute(0, 2, 3, 1).contiguous()
-        z_flattened = z.view(-1, self.e_dim)
-        ge_zero = (z_flattened > 0).float()
-        ones = torch.ones_like(z_flattened)
-        z_q = ones * ge_zero + -ones * (1 - ge_zero)
-
-        # preserve gradients
-        z_q = z_flattened + (z_q - z_flattened).detach()
-
-        # compute entropy loss
-        CatDist = torch.distributions.categorical.Categorical
-        logit = torch.stack(
-            [
-                -(z_flattened - torch.ones_like(z_q)).pow(2),
-                -(z_flattened - torch.ones_like(z_q) * -1).pow(2),
-            ],
-            dim=-1,
-        )
-        cat_dist = CatDist(logits=logit)
-        entropy = cat_dist.entropy().mean()
-        mean_prob = cat_dist.probs.mean(0)
-        mean_entropy = CatDist(probs=mean_prob).entropy().mean()
-
-        # compute loss for embedding
-        commit_loss = torch.mean(
-            (z_q.detach() - z_flattened) ** 2
-        ) + self.beta * torch.mean((z_q - z_flattened.detach()) ** 2)
-
-        # reshape back to match original input shape
-        z_q = z_q.view(z.shape)
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return {
-            "z": z_q,
-            "quantizer_loss": commit_loss * self.commit_loss_multiplier,
-            "entropy_loss": (entropy - mean_entropy) * self.entropy_multiplier,
-            "indices": self.get_indices(z_q),
-        }
-
-
-class VQGANDecoder(ModelMixin, ConfigMixin):
-    def __init__(self, ch: int = 128,
-                 ch_mult: List[int] = [1, 1, 2, 2, 4],
-                 num_res_blocks: List[int] = [4, 4, 3, 4, 3],
-                 attn_resolutions: List[int] = [5],
-                 dropout: float = 0.0,
-                 in_ch: int = 3,
-                 out_ch: int = 3,
-                 resolution: int = 256,
-                 z_channels: int = 13,
-                 double_z: bool = False):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_ch = in_ch
-        self.give_pre_end = False
-
-        self.z_channels = z_channels
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,) + tuple(ch_mult)
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        curr_res = self.resolution // 2 ** (self.num_resolutions - 1)
-        self.z_shape = (1, z_channels, curr_res, curr_res)
-        print(
-            "Working with z of shape {} = {} dimensions.".format(
-                self.z_shape, np.prod(self.z_shape)
-            )
-        )
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv2d(
-            z_channels, block_in, kernel_size=3, stride=1, padding=1
-        )
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(
-            in_channels=block_in,
-            out_channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-        )
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks[i_level]):
-                block.append(
-                    ResnetBlock(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        temb_channels=self.temb_ch,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, True)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(
-            block_in, out_ch, kernel_size=3, stride=1, padding=1
-        )
-        self.post_quant_conv = torch.nn.Conv2d(
-            z_channels, z_channels, 1
-        )
-
-
-    def forward(self, z):
-        # assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-        # timestep embedding
-        temb = None
-        output = dict()
-        z = self.post_quant_conv(z)
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks[i_level]):
-                h = self.up[i_level].block[i_block](h, temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        output["output"] = h
-        if self.give_pre_end:
-            return output
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        output["output"] = h
-        return output
-
-
-class MAGVITv2(ModelMixin, ConfigMixin):
-    @register_to_config
-    def __init__(
-            self,
-    ):
-        super().__init__()
-
-        self.encoder = VQGANEncoder()
-        self.decoder = VQGANDecoder()
-        self.quantize = LFQuantizer()
-
-    def forward(self, pixel_values, return_loss=False):
-        pass
-
-    def encode(self, pixel_values, return_loss=False):
-        hidden_states = self.encoder(pixel_values)
-        quantized_states = self.quantize(hidden_states)['z']
-        codebook_indices = self.quantize.get_indices(quantized_states).reshape(pixel_values.shape[0], -1)
-        output = (quantized_states, codebook_indices)
-        return output
-
-    def get_code(self, pixel_values):
-        hidden_states = self.encoder(pixel_values)
-        codebook_indices = self.quantize.get_indices(self.quantize(hidden_states)['z'])
-        codebook_indices = codebook_indices.reshape(pixel_values.shape[0], -1)
-
-        return codebook_indices
-
-    def decode_code(self, codebook_indices, shape=None):
-        z_q = self.quantize.get_codebook_entry(codebook_indices, shape=shape)
-
-        reconstructed_pixel_values = self.decoder(z_q)["output"]
-        return reconstructed_pixel_values
-
-
-if __name__ == '__main__':
-    encoder = VQGANEncoder()
-    import ipdb
-    ipdb.set_trace()
-    print()

omni_diffusion/models/magvit/modeling_utils.py

@@ -1,1184 +0,0 @@
-"""
-Modified from https://github.com/Gen-Verse/MMaDA/blob/main/models/modeling_utils.py
-"""
-
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team.
-# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-import itertools
-import json
-import os
-import re
-from collections import OrderedDict
-from functools import partial
-from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
-
-import safetensors
-import torch
-from huggingface_hub import create_repo, split_torch_state_dict_into_shards
-from huggingface_hub.utils import validate_hf_hub_args
-from torch import Tensor, nn
-
-from diffusers import __version__
-from diffusers.utils import (
-    FLAX_WEIGHTS_NAME,
-    SAFE_WEIGHTS_INDEX_NAME,
-    WEIGHTS_INDEX_NAME,
-    _add_variant,
-    _get_checkpoint_shard_files,
-    _get_model_file,
-    deprecate,
-    is_accelerate_available,
-    is_torch_version,
-    logging,
-)
-
-CONFIG_NAME = "config.json"
-WEIGHTS_NAME = "pytorch_model.bin"
-SAFETENSORS_WEIGHTS_NAME = "pytorch_model.safetensors"
-HUGGINGFACE_CO_RESOLVE_ENDPOINT = "https://huggingface.co"
-
-from diffusers.utils.hub_utils import (
-    PushToHubMixin,
-    load_or_create_model_card,
-    populate_model_card,
-)
-from diffusers.models.model_loading_utils import (
-    _determine_device_map,
-    _fetch_index_file,
-    _load_state_dict_into_model,
-    load_model_dict_into_meta,
-)
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-import diffusers
-
-logger = logging.get_logger(__name__)
-
-_REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}")
-
-
-if is_torch_version(">=", "1.9.0"):
-    _LOW_CPU_MEM_USAGE_DEFAULT = True
-else:
-    _LOW_CPU_MEM_USAGE_DEFAULT = False
-
-
-if is_accelerate_available():
-    import accelerate
-
-
-def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
-    try:
-        parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
-        return next(parameters_and_buffers).device
-    except StopIteration:
-        # For torch.nn.DataParallel compatibility in PyTorch 1.5
-
-        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
-            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
-            return tuples
-
-        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
-        first_tuple = next(gen)
-        return first_tuple[1].device
-
-
-def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
-    try:
-        params = tuple(parameter.parameters())
-        if len(params) > 0:
-            return params[0].dtype
-
-        buffers = tuple(parameter.buffers())
-        if len(buffers) > 0:
-            return buffers[0].dtype
-
-    except StopIteration:
-        # For torch.nn.DataParallel compatibility in PyTorch 1.5
-
-        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
-            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
-            return tuples
-
-        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
-        first_tuple = next(gen)
-        return first_tuple[1].dtype
-
-
-class ModelMixin(torch.nn.Module, PushToHubMixin):
-    r"""
-    Base class for all models.
-
-    [`ModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and
-    saving models.
-
-        - **config_name** ([`str`]) -- Filename to save a model to when calling [`~models.ModelMixin.save_pretrained`].
-    """
-
-    config_name = CONFIG_NAME
-    _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
-    _supports_gradient_checkpointing = False
-    _keys_to_ignore_on_load_unexpected = None
-    _no_split_modules = None
-
-    def __init__(self):
-        super().__init__()
-
-    def __getattr__(self, name: str) -> Any:
-        """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
-        config attributes directly. See https://github.com/huggingface/diffusers/pull/3129 We need to overwrite
-        __getattr__ here in addition so that we don't trigger `torch.nn.Module`'s __getattr__':
-        https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
-        """
-
-        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
-        is_attribute = name in self.__dict__
-
-        if is_in_config and not is_attribute:
-            deprecation_message = f"Accessing config attribute `{name}` directly via \
-                                '{type(self).__name__}' object attribute is deprecated. \
-                                Please access '{name}' over '{type(self).__name__}'s config object instead, \
-                                e.g. 'unet.config.{name}'."
-            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)
-            return self._internal_dict[name]
-
-        # call PyTorch's https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
-        return super().__getattr__(name)
-
-    @property
-    def is_gradient_checkpointing(self) -> bool:
-        """
-        Whether gradient checkpointing is activated for this model or not.
-        """
-        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
-
-    def enable_gradient_checkpointing(self) -> None:
-        """
-        Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
-        *checkpoint activations* in other frameworks).
-        """
-        if not self._supports_gradient_checkpointing:
-            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
-        self.apply(partial(self._set_gradient_checkpointing, value=True))
-
-    def disable_gradient_checkpointing(self) -> None:
-        """
-        Deactivates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
-        *checkpoint activations* in other frameworks).
-        """
-        if self._supports_gradient_checkpointing:
-            self.apply(partial(self._set_gradient_checkpointing, value=False))
-
-    # def set_use_npu_flash_attention(self, valid: bool) -> None:
-    #     r"""
-    #     Set the switch for the npu flash attention.
-    #     """
-
-    #     def fn_recursive_set_npu_flash_attention(module: torch.nn.Module):
-    #         if hasattr(module, "set_use_npu_flash_attention"):
-    #             module.set_use_npu_flash_attention(valid)
-
-    #         for child in module.children():
-    #             fn_recursive_set_npu_flash_attention(child)
-
-    #     for module in self.children():
-    #         if isinstance(module, torch.nn.Module):
-    #             fn_recursive_set_npu_flash_attention(module)
-
-    # def enable_npu_flash_attention(self) -> None:
-    #     r"""
-    #     Enable npu flash attention from torch_npu
-
-    #     """
-    #     self.set_use_npu_flash_attention(True)
-
-    # def disable_npu_flash_attention(self) -> None:
-    #     r"""
-    #     disable npu flash attention from torch_npu
-
-    #     """
-    #     self.set_use_npu_flash_attention(False)
-
-    # def set_use_memory_efficient_attention_xformers(
-    #     self, valid: bool, attention_op: Optional[Callable] = None
-    # ) -> None:
-    #     # Recursively walk through all the children.
-    #     # Any children which exposes the set_use_memory_efficient_attention_xformers method
-    #     # gets the message
-    #     def fn_recursive_set_mem_eff(module: torch.nn.Module):
-    #         if hasattr(module, "set_use_memory_efficient_attention_xformers"):
-    #             module.set_use_memory_efficient_attention_xformers(valid, attention_op)
-
-    #         for child in module.children():
-    #             fn_recursive_set_mem_eff(child)
-
-    #     for module in self.children():
-    #         if isinstance(module, torch.nn.Module):
-    #             fn_recursive_set_mem_eff(module)
-
-    # def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None) -> None:
-    #     r"""
-    #     Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
-
-    #     When this option is enabled, you should observe lower GPU memory usage and a potential speed up during
-    #     inference. Speed up during training is not guaranteed.
-
-    #     <Tip warning={true}>
-
-    #     ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes
-    #     precedent.
-
-    #     </Tip>
-
-    #     Parameters:
-    #         attention_op (`Callable`, *optional*):
-    #             Override the default `None` operator for use as `op` argument to the
-    #             [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#x
-    #             formers.ops.memory_efficient_attention)
-    #             function of xFormers.
-
-    #     Examples:
-
-    #     ```py
-    #     >>> import torch
-    #     >>> from diffusers import UNet2DConditionModel
-    #     >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
-
-    #     >>> model = UNet2DConditionModel.from_pretrained(
-    #     ...     "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
-    #     ... )
-    #     >>> model = model.to("cuda")
-    #     >>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
-    #     ```
-    #     """
-    #     self.set_use_memory_efficient_attention_xformers(True, attention_op)
-
-    # def disable_xformers_memory_efficient_attention(self) -> None:
-    #     r"""
-    #     Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
-    #     """
-    #     self.set_use_memory_efficient_attention_xformers(False)
-
-    def save_pretrained(
-        self,
-        save_directory: Union[str, os.PathLike],
-        is_main_process: bool = True,
-        save_function: Optional[Callable] = None,
-        safe_serialization: bool = True,
-        variant: Optional[str] = None,
-        max_shard_size: Union[int, str] = "10GB",
-        push_to_hub: bool = False,
-        **kwargs,
-    ):
-        """
-        Save a model and its configuration file to a directory so that it can be reloaded using the
-        [`~models.ModelMixin.from_pretrained`] class method.
-
-        Arguments:
-            save_directory (`str` or `os.PathLike`):
-                Directory to save a model and its configuration file to. Will be created if it doesn't exist.
-            is_main_process (`bool`, *optional*, defaults to `True`):
-                Whether the process calling this is the main process or not. Useful during distributed training and you
-                need to call this function on all processes. In this case, set `is_main_process=True` only on the main
-                process to avoid race conditions.
-            save_function (`Callable`):
-                The function to use to save the state dictionary. Useful during distributed training when you need to
-                replace `torch.save` with another method. Can be configured with the environment variable
-                `DIFFUSERS_SAVE_MODE`.
-            safe_serialization (`bool`, *optional*, defaults to `True`):
-                Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
-            variant (`str`, *optional*):
-                If specified, weights are saved in the format `pytorch_model.<variant>.bin`.
-            max_shard_size (`int` or `str`, defaults to `"10GB"`):
-                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
-                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5GB"`).
-                If expressed as an integer, the unit is bytes. Note that this limit will be decreased after a certain
-                period of time (starting from Oct 2024) to allow users to upgrade to the latest version of `diffusers`.
-                This is to establish a common default size for this argument across different libraries in the Hugging
-                Face ecosystem (`transformers`, and `accelerate`, for example).
-            push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
-                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
-                namespace).
-            kwargs (`Dict[str, Any]`, *optional*):
-                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
-        """
-        if os.path.isfile(save_directory):
-            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
-            return
-
-        weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
-        weights_name = _add_variant(weights_name, variant)
-        weight_name_split = weights_name.split(".")
-        if len(weight_name_split) in [2, 3]:
-            weights_name_pattern = weight_name_split[0] + "{suffix}." + ".".join(weight_name_split[1:])
-        else:
-            raise ValueError(f"Invalid {weights_name} provided.")
-
-        os.makedirs(save_directory, exist_ok=True)
-
-        if push_to_hub:
-            commit_message = kwargs.pop("commit_message", None)
-            private = kwargs.pop("private", False)
-            create_pr = kwargs.pop("create_pr", False)
-            token = kwargs.pop("token", None)
-            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
-            repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
-
-        # Only save the model itself if we are using distributed training
-        model_to_save = self
-
-        # Attach architecture to the config
-        # Save the config
-        if is_main_process:
-            model_to_save.save_config(save_directory)
-
-        # Save the model
-        state_dict = model_to_save.state_dict()
-
-        # Save the model
-        state_dict_split = split_torch_state_dict_into_shards(
-            state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern
-        )
-
-        # Clean the folder from a previous save
-        if is_main_process:
-            for filename in os.listdir(save_directory):
-                if filename in state_dict_split.filename_to_tensors.keys():
-                    continue
-                full_filename = os.path.join(save_directory, filename)
-                if not os.path.isfile(full_filename):
-                    continue
-                weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "")
-                weights_without_ext = weights_without_ext.replace("{suffix}", "")
-                filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "")
-                # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
-                if (
-                    filename.startswith(weights_without_ext)
-                    and _REGEX_SHARD.fullmatch(filename_without_ext) is not None
-                ):
-                    os.remove(full_filename)
-
-        for filename, tensors in state_dict_split.filename_to_tensors.items():
-            shard = {tensor: state_dict[tensor] for tensor in tensors}
-            filepath = os.path.join(save_directory, filename)
-            if safe_serialization:
-                # At some point we will need to deal better with save_function (used for TPU and other distributed
-                # joyfulness), but for now this enough.
-                safetensors.torch.save_file(shard, filepath, metadata={"format": "pt"})
-            else:
-                torch.save(shard, filepath)
-
-        if state_dict_split.is_sharded:
-            index = {
-                "metadata": state_dict_split.metadata,
-                "weight_map": state_dict_split.tensor_to_filename,
-            }
-            save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
-            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-            # Save the index as well
-            with open(save_index_file, "w", encoding="utf-8") as f:
-                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                f.write(content)
-            logger.info(
-                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each "
-                f"parameters has been saved in the index located at {save_index_file}."
-            )
-        else:
-            path_to_weights = os.path.join(save_directory, weights_name)
-            logger.info(f"Model weights saved in {path_to_weights}")
-
-        if push_to_hub:
-            # Create a new empty model card and eventually tag it
-            model_card = load_or_create_model_card(repo_id, token=token)
-            model_card = populate_model_card(model_card)
-            model_card.save(Path(save_directory, "README.md").as_posix())
-
-            self._upload_folder(
-                save_directory,
-                repo_id,
-                token=token,
-                commit_message=commit_message,
-                create_pr=create_pr,
-            )
-
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        r"""
-        Instantiate a pretrained PyTorch model from a pretrained model configuration.
-
-        The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To
-        train the model, set it back in training mode with `model.train()`.
-
-        Parameters:
-            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
-                Can be either:
-
-                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
-                      the Hub.
-                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
-                      with [`~ModelMixin.save_pretrained`].
-
-            cache_dir (`Union[str, os.PathLike]`, *optional*):
-                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
-                is not used.
-            torch_dtype (`str` or `torch.dtype`, *optional*):
-                Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
-                dtype is automatically derived from the model's weights.
-            force_download (`bool`, *optional*, defaults to `False`):
-                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
-                cached versions if they exist.
-            proxies (`Dict[str, str]`, *optional*):
-                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
-                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
-            output_loading_info (`bool`, *optional*, defaults to `False`):
-                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
-            local_files_only(`bool`, *optional*, defaults to `False`):
-                Whether to only load local model weights and configuration files or not. If set to `True`, the model
-                won't be downloaded from the Hub.
-            token (`str` or *bool*, *optional*):
-                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
-                `diffusers-cli login` (stored in `~/.huggingface`) is used.
-            revision (`str`, *optional*, defaults to `"main"`):
-                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
-                allowed by Git.
-            from_flax (`bool`, *optional*, defaults to `False`):
-                Load the model weights from a Flax checkpoint save file.
-            subfolder (`str`, *optional*, defaults to `""`):
-                The subfolder location of a model file within a larger model repository on the Hub or locally.
-            mirror (`str`, *optional*):
-                Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
-                guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
-                information.
-            device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
-                A map that specifies where each submodule should go. It doesn't need to be defined for each
-                parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
-                same device. Defaults to `None`, meaning that the model will be loaded on CPU.
-
-                Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For
-                more information about each option see [designing a device
-                map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
-            max_memory (`Dict`, *optional*):
-                A dictionary device identifier for the maximum memory. Will default to the maximum memory available for
-                each GPU and the available CPU RAM if unset.
-            offload_folder (`str` or `os.PathLike`, *optional*):
-                The path to offload weights if `device_map` contains the value `"disk"`.
-            offload_state_dict (`bool`, *optional*):
-                If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if
-                the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True`
-                when there is some disk offload.
-            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
-                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
-                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
-                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
-                argument to `True` will raise an error.
-            variant (`str`, *optional*):
-                Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when
-                loading `from_flax`.
-            use_safetensors (`bool`, *optional*, defaults to `None`):
-                If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the
-                `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors`
-                weights. If set to `False`, `safetensors` weights are not loaded.
-
-        <Tip>
-
-        To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with
-        `huggingface-cli login`. You can also activate the special
-        ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a
-        firewalled environment.
-
-        </Tip>
-
-        Example:
-
-        ```py
-        from diffusers import UNet2DConditionModel
-
-        unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
-        ```
-
-        If you get the error message below, you need to finetune the weights for your downstream task:
-
-        ```bash
-        Some weights of UNet2DConditionModel were not initialized from the model checkpoint at 
-        runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match:
-        - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) 
-            in the model instantiated
-        You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
-        ```
-        """
-        cache_dir = kwargs.pop("cache_dir", None)
-        ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
-        force_download = kwargs.pop("force_download", False)
-        from_flax = kwargs.pop("from_flax", False)
-        proxies = kwargs.pop("proxies", None)
-        output_loading_info = kwargs.pop("output_loading_info", False)
-        local_files_only = kwargs.pop("local_files_only", None)
-        token = kwargs.pop("token", None)
-        revision = kwargs.pop("revision", None)
-        torch_dtype = kwargs.pop("torch_dtype", None)
-        subfolder = kwargs.pop("subfolder", None)
-        device_map = kwargs.pop("device_map", None)
-        max_memory = kwargs.pop("max_memory", None)
-        offload_folder = kwargs.pop("offload_folder", None)
-        offload_state_dict = kwargs.pop("offload_state_dict", False)
-        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
-        variant = kwargs.pop("variant", None)
-        use_safetensors = kwargs.pop("use_safetensors", None)
-
-        allow_pickle = False
-        if use_safetensors is None:
-            use_safetensors = True
-            allow_pickle = True
-
-        def _check_input(low_cpu_mem_usage, device_map):
-            if low_cpu_mem_usage and not is_accelerate_available():
-                low_cpu_mem_usage = False
-                logger.warning(
-                    "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
-                    " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
-                    " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
-                    " install accelerate\n```\n."
-                )
-
-            if device_map is not None and not is_accelerate_available():
-                raise NotImplementedError(
-                    "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set"
-                    " `device_map=None`. You can install accelerate with `pip install accelerate`."
-                )
-
-            # Check if we can handle device_map and dispatching the weights
-            if device_map is not None and not is_torch_version(">=", "1.9.0"):
-                raise NotImplementedError(
-                    "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
-                    " `device_map=None`."
-                )
-
-            if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
-                raise NotImplementedError(
-                    "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or "
-                    "set `low_cpu_mem_usage=False`."
-                )
-
-            if low_cpu_mem_usage is False and device_map is not None:
-                raise ValueError(
-                    f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and"
-                    " dispatching. Please make sure to set `low_cpu_mem_usage=True`."
-                )
-        _check_input(low_cpu_mem_usage, device_map)
-
-        # change device_map into a map if we passed an int, a str or a torch.device
-        if isinstance(device_map, torch.device):
-            device_map = {"": device_map}
-        elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
-            try:
-                device_map = {"": torch.device(device_map)}
-            except RuntimeError:
-                raise ValueError(
-                    "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or "
-                    f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}."
-                )
-        elif isinstance(device_map, int):
-            if device_map < 0:
-                raise ValueError(
-                    "You can't pass device_map as a negative int. " \
-                    "If you want to put the model on the cpu, pass device_map = 'cpu' "
-                )
-            else:
-                device_map = {"": device_map}
-
-        if device_map is not None:
-            if low_cpu_mem_usage is None:
-                low_cpu_mem_usage = True
-            elif not low_cpu_mem_usage:
-                raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`")
-
-        if low_cpu_mem_usage:
-            if device_map is not None and not is_torch_version(">=", "1.10"):
-                # The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info.
-                raise ValueError("`low_cpu_mem_usage` and `device_map` require PyTorch >= 1.10.")
-
-        # Load config if we don't provide a configuration
-        config_path = pretrained_model_name_or_path
-
-        user_agent = {
-            "diffusers": __version__,
-            "file_type": "model",
-            "framework": "pytorch",
-        }
-
-        # load config
-        def _load_config(cls, config_path, **kwargs):
-            config, unused_kwargs, commit_hash = cls.load_config(
-                config_path,
-                cache_dir=cache_dir,
-                return_unused_kwargs=True,
-                return_commit_hash=True,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                subfolder=subfolder,
-                user_agent=user_agent,
-                **kwargs,
-            )
-
-            # Determine if we're loading from a directory of sharded checkpoints.
-            is_sharded = False
-            index_file = None
-            is_local = os.path.isdir(pretrained_model_name_or_path)
-            index_file = _fetch_index_file(
-                is_local=is_local,
-                pretrained_model_name_or_path=pretrained_model_name_or_path,
-                subfolder=subfolder or "",
-                use_safetensors=use_safetensors,
-                cache_dir=cache_dir,
-                variant=variant,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                user_agent=user_agent,
-                commit_hash=commit_hash,
-            )
-            if index_file is not None and index_file.is_file():
-                is_sharded = True
-
-            if is_sharded and from_flax:
-                raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.")
-            return config, unused_kwargs, commit_hash, is_sharded, index_file
-        
-        config, unused_kwargs, commit_hash, is_sharded, index_file = _load_config(cls, config_path, **kwargs)
-
-        # load model
-        model_file = None
-        if from_flax:
-            model_file = _get_model_file(
-                pretrained_model_name_or_path,
-                weights_name=FLAX_WEIGHTS_NAME,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                local_files_only=local_files_only,
-                token=token,
-                revision=revision,
-                subfolder=subfolder,
-                user_agent=user_agent,
-                commit_hash=commit_hash,
-            )
-            model = cls.from_config(config, **unused_kwargs)
-
-            # Convert the weights
-            from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model
-
-            model = load_flax_checkpoint_in_pytorch_model(model, model_file)
-        else:
-            if is_sharded:
-                sharded_ckpt_cached_folder, sharded_metadata = _get_checkpoint_shard_files(
-                    pretrained_model_name_or_path,
-                    index_file,
-                    cache_dir=cache_dir,
-                    proxies=proxies,
-                    local_files_only=local_files_only,
-                    token=token,
-                    user_agent=user_agent,
-                    revision=revision,
-                    subfolder=subfolder or "",
-                )
-
-            elif use_safetensors and not is_sharded:
-                try:
-                    model_file = _get_model_file(
-                        pretrained_model_name_or_path,
-                        weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
-                        cache_dir=cache_dir,
-                        force_download=force_download,
-                        proxies=proxies,
-                        local_files_only=local_files_only,
-                        token=token,
-                        revision=revision,
-                        subfolder=subfolder,
-                        user_agent=user_agent,
-                        commit_hash=commit_hash,
-                    )
-
-                except IOError as e:
-                    logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
-                    if not allow_pickle:
-                        raise
-                    logger.warning(
-                        "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
-                    )
-
-            if model_file is None and not is_sharded:
-                model_file = _get_model_file(
-                    pretrained_model_name_or_path,
-                    weights_name=_add_variant(WEIGHTS_NAME, variant),
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    proxies=proxies,
-                    local_files_only=local_files_only,
-                    token=token,
-                    revision=revision,
-                    subfolder=subfolder,
-                    user_agent=user_agent,
-                    commit_hash=commit_hash,
-                )
-
-            if low_cpu_mem_usage:
-                # Instantiate model with empty weights
-                def load_model_in_low_cpu_mem_usage(config, device_map, unused_kwargs, commit_hash, is_sharded, 
-                    index_file):
-                    with accelerate.init_empty_weights():
-                        model = cls.from_config(config, **unused_kwargs)
-
-                    # if device_map is None, load the state dict and move the params from meta device to the cpu
-                    if device_map is None and not is_sharded:
-                        param_device = "cpu"
-                        state_dict = diffusers.models.model_loading_utils.load_state_dict(model_file, variant=variant)
-                        model._convert_deprecated_attn(state_dict)
-                        # move the params from meta device to cpu
-                        missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
-                        if len(missing_keys) > 0:
-                            raise ValueError(
-                                f"Cannot load {cls} from {pretrained_model_name_or_path} because the following keys are"
-                                f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
-                                " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
-                                " those weights or else make sure your checkpoint file is correct."
-                            )
-
-                        unexpected_keys = load_model_dict_into_meta(
-                            model,
-                            state_dict,
-                            device=param_device,
-                            dtype=torch_dtype,
-                            model_name_or_path=pretrained_model_name_or_path,
-                        )
-
-                        if cls._keys_to_ignore_on_load_unexpected is not None:
-                            for pat in cls._keys_to_ignore_on_load_unexpected:
-                                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
-
-                        if len(unexpected_keys) > 0:
-                            logger.warning(
-                                f"Some weights of the model checkpoint were not used when initializing "
-                                f" {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
-                            )
-
-                    else:  # else let accelerate handle loading and dispatching.
-                        # Load weights and dispatch according to the device_map
-                        # by default the device_map is None and the weights are loaded on the CPU
-                        force_hook = True
-                        device_map = _determine_device_map(model, device_map, max_memory, torch_dtype)
-                        if device_map is None and is_sharded:
-                            # we load the parameters on the cpu
-                            device_map = {"": "cpu"}
-                            force_hook = False
-                        try:
-                            accelerate.load_checkpoint_and_dispatch(
-                                model,
-                                model_file if not is_sharded else index_file,
-                                device_map,
-                                max_memory=max_memory,
-                                offload_folder=offload_folder,
-                                offload_state_dict=offload_state_dict,
-                                dtype=torch_dtype,
-                                force_hooks=force_hook,
-                                strict=True,
-                            )
-                        except AttributeError as e:
-                            # When using accelerate loading, we do not have the ability to load the state
-                            # dict and rename the weight names manually. Additionally, accelerate skips
-                            # torch loading conventions and directly writes into `module.{_buffers, _parameters}`
-                            # (which look like they should be private variables?), so we can't use the standard hooks
-                            # to rename parameters on load. We need to mimic the original weight names so the correct
-                            # attributes are available. After we have loaded the weights, we convert the deprecated
-                            # names to the new non-deprecated names. Then we _greatly encourage_ the user to convert
-                            # the weights so we don't have to do this again.
-
-                            if "'Attention' object has no attribute" in str(e):
-                                logger.warning(
-                                    f"Taking `{str(e)}` while using `accelerate.load_checkpoint_and_dispatch` to mean "
-                                    f"{pretrained_model_name_or_path} was saved with deprecated attention block weight "
-                                    f"names. We will load it with the deprecated attention block names and convert them"
-                                    f" on the fly to the new attention block format. Please re-save the model after "
-                                    f"this conversion so we don't have to do the on the fly renaming in the future. "
-                                    f"If the model is from a hub checkpoint, please also re-upload it or open a PR on"
-                                    f" the original repository."
-                                )
-                                model._temp_convert_self_to_deprecated_attention_blocks()
-                                accelerate.load_checkpoint_and_dispatch(
-                                    model,
-                                    model_file if not is_sharded else index_file,
-                                    device_map,
-                                    max_memory=max_memory,
-                                    offload_folder=offload_folder,
-                                    offload_state_dict=offload_state_dict,
-                                    dtype=torch_dtype,
-                                    force_hooks=force_hook,
-                                    strict=True,
-                                )
-                                model._undo_temp_convert_self_to_deprecated_attention_blocks()
-                            else:
-                                raise e
-
-                    loading_info = {
-                        "missing_keys": [],
-                        "unexpected_keys": [],
-                        "mismatched_keys": [],
-                        "error_msgs": [],
-                    }
-                    return model, loading_info
-                model, loading_info = load_model_in_low_cpu_mem_usage(config, device_map, unused_kwargs, commit_hash, 
-                                                                      is_sharded, index_file)
-
-            else:
-                model = cls.from_config(config, **unused_kwargs)
-
-                state_dict = load_state_dict(model_file, variant=variant)
-                model._convert_deprecated_attn(state_dict)
-
-                model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
-                    model,
-                    state_dict,
-                    model_file,
-                    pretrained_model_name_or_path,
-                    ignore_mismatched_sizes=ignore_mismatched_sizes,
-                )
-
-                loading_info = {
-                    "missing_keys": missing_keys,
-                    "unexpected_keys": unexpected_keys,
-                    "mismatched_keys": mismatched_keys,
-                    "error_msgs": error_msgs,
-                }
-
-        if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
-            raise ValueError(
-                f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
-            )
-        elif torch_dtype is not None:
-            model = model.to(torch_dtype)
-    
-        model.register_to_config(_name_or_path=pretrained_model_name_or_path)
-
-        # Set model in evaluation mode to deactivate DropOut modules by default
-        model.eval()
-        if output_loading_info:
-            return model, loading_info
-
-        return model
-
-    @classmethod
-    def _load_pretrained_model(
-        cls,
-        model,
-        state_dict: OrderedDict,
-        resolved_archive_file,
-        pretrained_model_name_or_path: Union[str, os.PathLike],
-        ignore_mismatched_sizes: bool = False,
-    ):
-        # Retrieve missing & unexpected_keys
-        model_state_dict = model.state_dict()
-        loaded_keys = list(state_dict.keys())
-
-        expected_keys = list(model_state_dict.keys())
-
-        original_loaded_keys = loaded_keys
-
-        missing_keys = list(set(expected_keys) - set(loaded_keys))
-        unexpected_keys = list(set(loaded_keys) - set(expected_keys))
-
-        # Make sure we are able to load base models as well as derived models (with heads)
-        model_to_load = model
-
-        def _find_mismatched_keys(
-            state_dict,
-            model_state_dict,
-            loaded_keys,
-            ignore_mismatched_sizes,
-        ):
-            mismatched_keys = []
-            if ignore_mismatched_sizes:
-                for checkpoint_key in loaded_keys:
-                    model_key = checkpoint_key
-
-                    if (
-                        model_key in model_state_dict
-                        and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
-                    ):
-                        mismatched_keys.append(
-                            (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
-                        )
-                        del state_dict[checkpoint_key]
-            return mismatched_keys
-
-        if state_dict is not None:
-            # Whole checkpoint
-            mismatched_keys = _find_mismatched_keys(
-                state_dict,
-                model_state_dict,
-                original_loaded_keys,
-                ignore_mismatched_sizes,
-            )
-            error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
-
-        if len(error_msgs) > 0:
-            error_msg = "\n\t".join(error_msgs)
-            if "size mismatch" in error_msg:
-                error_msg += (
-                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
-                )
-            raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
-
-        if len(unexpected_keys) > 0:
-            logger.warning(
-                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
-                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
-                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task"
-                " or with another architecture (e.g. initializing a BertForSequenceClassification model from a"
-                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
-                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly"
-                " identical (initializing a BertForSequenceClassification model from a"
-                " BertForSequenceClassification model)."
-            )
-        else:
-            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
-        if len(missing_keys) > 0:
-            logger.warning(
-                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
-                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
-            )
-        elif len(mismatched_keys) == 0:
-            logger.info(
-                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the"
-                f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions"
-                " without further training."
-            )
-        if len(mismatched_keys) > 0:
-            mismatched_warning = "\n".join(
-                [
-                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
-                    for key, shape1, shape2 in mismatched_keys
-                ]
-            )
-            logger.warning(
-                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
-                f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
-                f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be"
-                " able to use it for predictions and inference."
-            )
-
-        return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
-
-    @classmethod
-    def _get_signature_keys(cls, obj):
-        parameters = inspect.signature(obj.__init__).parameters
-        required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
-        optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
-        expected_modules = set(required_parameters.keys()) - {"self"}
-
-        return expected_modules, optional_parameters
-
-    # Adapted from `transformers` modeling_utils.py
-    def _get_no_split_modules(self, device_map: str):
-        """
-        Get the modules of the model that should not be spit when using device_map. We iterate through the modules to
-        get the underlying `_no_split_modules`.
-
-        Args:
-            device_map (`str`):
-                The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"]
-
-        Returns:
-            `List[str]`: List of modules that should not be split
-        """
-        _no_split_modules = set()
-        modules_to_check = [self]
-        while len(modules_to_check) > 0:
-            module = modules_to_check.pop(-1)
-            # if the module does not appear in _no_split_modules, we also check the children
-            if module.__class__.__name__ not in _no_split_modules:
-                if isinstance(module, ModelMixin):
-                    if module._no_split_modules is None:
-                        raise ValueError(
-                            f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement "
-                            f"support, the model class needs to implement the `_no_split_modules` attribute."
-                        )
-                    else:
-                        _no_split_modules = _no_split_modules | set(module._no_split_modules)
-                modules_to_check += list(module.children())
-        return list(_no_split_modules)
-
-    @property
-    def device(self) -> torch.device:
-        """
-        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
-        device).
-        """
-        return get_parameter_device(self)
-
-    @property
-    def dtype(self) -> torch.dtype:
-        """
-        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
-        """
-        return get_parameter_dtype(self)
-
-    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
-        """
-        Get number of (trainable or non-embedding) parameters in the module.
-
-        Args:
-            only_trainable (`bool`, *optional*, defaults to `False`):
-                Whether or not to return only the number of trainable parameters.
-            exclude_embeddings (`bool`, *optional*, defaults to `False`):
-                Whether or not to return only the number of non-embedding parameters.
-
-        Returns:
-            `int`: The number of parameters.
-
-        Example:
-
-        ```py
-        from diffusers import UNet2DConditionModel
-
-        model_id = "runwayml/stable-diffusion-v1-5"
-        unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
-        unet.num_parameters(only_trainable=True)
-        859520964
-        ```
-        """
-
-        if exclude_embeddings:
-            embedding_param_names = [
-                f"{name}.weight"
-                for name, module_type in self.named_modules()
-                if isinstance(module_type, torch.nn.Embedding)
-            ]
-            non_embedding_parameters = [
-                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
-            ]
-            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
-        else:
-            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
-
-    def _convert_deprecated_attn(self, state_dict: OrderedDict) -> None:
-        # original function name: _convert_deprecated_attention_blocks
-        deprecated_attention_block_paths = []
-
-        def recursive_find_attn_block(name, module):
-            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
-                deprecated_attention_block_paths.append(name)
-
-            for sub_name, sub_module in module.named_children():
-                sub_name = sub_name if name == "" else f"{name}.{sub_name}"
-                recursive_find_attn_block(sub_name, sub_module)
-
-        recursive_find_attn_block("", self)
-
-        # NOTE: we have to check if the deprecated parameters are in the state dict
-        # because it is possible we are loading from a state dict that was already
-        # converted
-
-        for path in deprecated_attention_block_paths:
-            # group_norm path stays the same
-
-            # query -> to_q
-            if f"{path}.query.weight" in state_dict:
-                state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")
-            if f"{path}.query.bias" in state_dict:
-                state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")
-
-            # key -> to_k
-            if f"{path}.key.weight" in state_dict:
-                state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")
-            if f"{path}.key.bias" in state_dict:
-                state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")
-
-            # value -> to_v
-            if f"{path}.value.weight" in state_dict:
-                state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")
-            if f"{path}.value.bias" in state_dict:
-                state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")
-
-            # proj_attn -> to_out.0
-            if f"{path}.proj_attn.weight" in state_dict:
-                state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")
-            if f"{path}.proj_attn.bias" in state_dict:
-                state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")
-
-
-class LegacyModelMixin(ModelMixin):
-    r"""
-    A subclass of `ModelMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
-    pipeline-specific classes (like `DiTTransformer2DModel`).
-    """
-
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        # To prevent dependency import problem.
-        from diffusers.models.model_loading_utils import _fetch_remapped_cls_from_config
-
-        # Create a copy of the kwargs so that we don't mess with the keyword arguments in the downstream calls.
-        kwargs_copy = kwargs.copy()
-
-        cache_dir = kwargs.pop("cache_dir", None)
-        force_download = kwargs.pop("force_download", False)
-        proxies = kwargs.pop("proxies", None)
-        local_files_only = kwargs.pop("local_files_only", None)
-        token = kwargs.pop("token", None)
-        revision = kwargs.pop("revision", None)
-        subfolder = kwargs.pop("subfolder", None)
-
-        # Load config if we don't provide a configuration
-        config_path = pretrained_model_name_or_path
-
-        user_agent = {
-            "diffusers": __version__,
-            "file_type": "model",
-            "framework": "pytorch",
-        }
-
-        # load config
-        config, _, _ = cls.load_config(
-            config_path,
-            cache_dir=cache_dir,
-            return_unused_kwargs=True,
-            return_commit_hash=True,
-            force_download=force_download,
-            proxies=proxies,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            subfolder=subfolder,
-            user_agent=user_agent,
-            **kwargs,
-        )
-        # resolve remapping
-        remapped_class = _fetch_remapped_cls_from_config(config, cls)
-
-        return remapped_class.from_pretrained(pretrained_model_name_or_path, **kwargs_copy)

omni_diffusion/tokenizer.py

@@ -1,74 +0,0 @@
-from .constants import (
-    BOX_END_TOKEN,
-    BOX_START_TOKEN,
-    IMG_CONTEXT_TOKEN,
-    IMG_END_TOKEN,
-    IMG_START_TOKEN,
-    IMG_TAG_TOKEN,
-    PATCH_CONTEXT_TOKEN,
-    PATCH_END_TOKEN,
-    PATCH_START_TOKEN,
-    QUAD_END_TOKEN,
-    QUAD_START_TOKEN,
-    REF_END_TOKEN,
-    REF_START_TOKEN,
-    VID_CONTEXT_TOKEN,
-    VID_END_TOKEN,
-    VID_START_TOKEN,
-    VID_TAG_TOKEN,
-)
-
-
-def update_tokenizer(tokenizer, audio_tokenizer_type=None):
-    """
-    Adds special tokens related to multimodal tasks (Image, Video, Patch, Bounding Box, etc.)
-    to the tokenizer vocabulary.
-    Also handles specific tokenizer updates for audio models like SenseVoice/GLM4Voice.
-    """
-    token_list = [
-        IMG_START_TOKEN,
-        IMG_END_TOKEN,
-        IMG_CONTEXT_TOKEN,
-        VID_START_TOKEN,
-        VID_END_TOKEN,
-        VID_CONTEXT_TOKEN,
-        PATCH_START_TOKEN,
-        PATCH_END_TOKEN,
-        PATCH_CONTEXT_TOKEN,
-        QUAD_START_TOKEN,
-        QUAD_END_TOKEN,
-        REF_START_TOKEN,
-        REF_END_TOKEN,
-        BOX_START_TOKEN,
-        BOX_END_TOKEN,
-        IMG_TAG_TOKEN,
-        VID_TAG_TOKEN,
-    ]
-    num_new_tokens = tokenizer.add_tokens(token_list, special_tokens=True)
-
-    if audio_tokenizer_type == "sensevoice_glm4voice":
-        from .tokenizer_sensevoice_glm4voice import (
-            update_tokenizer_for_sensevoice_glm4voice,
-            SenseVoiceGLM4VoiceTokenizer,
-        )
-        return update_tokenizer_for_sensevoice_glm4voice(tokenizer)
-
-    raise NotImplementedError
-
-
-def get_audio_tokenizer(model_name_or_path, audio_tokenizer_type, flow_path=None, rank=None):
-    """
-    Factory function to initialize and return the appropriate audio tokenizer instance
-    based on the provided type (e.g., 'sensevoice_glm4voice').
-    """
-    if audio_tokenizer_type is None:
-        return None
-
-    if audio_tokenizer_type == "sensevoice_glm4voice":
-        from .tokenizer_sensevoice_glm4voice import (
-            update_tokenizer_for_sensevoice_glm4voice,
-            SenseVoiceGLM4VoiceTokenizer,
-        )
-        return SenseVoiceGLM4VoiceTokenizer(model_name_or_path, flow_path=flow_path, rank=rank)
-
-    raise NotImplementedError

omni_diffusion/tokenizer_magvitv2.py

@@ -1,65 +0,0 @@
-import glob
-import io
-import logging
-import math
-import os
-import tarfile
-import uuid
-
-import safetensors
-import torch
-from transformers import WhisperFeatureExtractor, WhisperTokenizerFast
-
-import torchaudio
-
-from transformers import WhisperFeatureExtractor
-from speech_tokenizer.modeling_whisper import WhisperVQEncoder
-from flow_inference import AudioDecoder
-from omni_diffusion.models.magvit.modeling_magvitv2 import MAGVITv2
-
-from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
-from funasr.models.sense_voice.model import SenseVoiceSmall
-            
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-class MagVITV2Tokenizer:
-    def __init__(self, model_path=None, rank=None):
-        if rank is None and torch.distributed.is_initialized():
-            rank = torch.distributed.get_rank()
-            self.rank = rank % 8
-        else:
-            self.rank = rank
-        logger.info(f"{self.rank=}")
-        self.model_path = model_path
-
-    def load_model(self):
-
-        if hasattr(self, "image_tokenizer"):
-            return
-
-        if self.rank is not None:
-            self.device = f"cuda:{self.rank}"
-            torch.cuda.set_device(self.rank)
-        else:
-            self.device = "cpu"
-
-        self.image_tokenizer = MAGVITv2()
-        logger.info(f"{self.device=} Loading MAGVITv2")
-        self.image_tokenizer = self.image_tokenizer.from_pretrained(self.model_path).to(self.device)
-        self.image_tokenizer.eval()
-        self.image_tokenizer.requires_grad_(False) 
-        logger.info(f"{self.device=} Loading MAGVITv2 done")
-
-    def encode(self, image):
-        image = image.to(self.device)
-        image_tokens = self.image_tokenizer.get_code(image)
-        return image_tokens
-
-    def decode(self, image_tokens):
-        images = self.image_tokenizer.decode_code(image_tokens)
-        return images
-
-    def apply_to_role(self, role, **kwargs):
-        return True

omni_diffusion/tokenizer_sensevoice_glm4voice.py

@@ -1,280 +0,0 @@
-import glob
-import io
-import logging
-import math
-import os
-import tarfile
-import uuid
-
-import safetensors
-import torch
-from transformers import WhisperFeatureExtractor, WhisperTokenizerFast
-
-import torchaudio
-
-from transformers import WhisperFeatureExtractor
-from speech_tokenizer.modeling_whisper import WhisperVQEncoder
-from flow_inference import AudioDecoder
-
-from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
-from funasr.models.sense_voice.model import SenseVoiceSmall
-import sys
-            
-from .constants import (
-    AUD_CONTEXT_TOKEN,
-    AUD_END_TOKEN,
-    AUD_START_TOKEN,
-    AUD_TAG_TOKEN,
-    BOX_END_TOKEN,
-    BOX_START_TOKEN,
-    IMG_CONTEXT_TOKEN,
-    IMG_END_TOKEN,
-    IMG_START_TOKEN,
-    IMG_TAG_TOKEN,
-    PATCH_CONTEXT_TOKEN,
-    PATCH_END_TOKEN,
-    PATCH_START_TOKEN,
-    QUAD_END_TOKEN,
-    QUAD_START_TOKEN,
-    REF_END_TOKEN,
-    REF_START_TOKEN,
-    VID_CONTEXT_TOKEN,
-    VID_END_TOKEN,
-    VID_START_TOKEN,
-    VID_TAG_TOKEN,
-)
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-def update_tokenizer_for_sensevoice_glm4voice(tokenizer):
-    token_list = [
-        IMG_START_TOKEN,
-        IMG_END_TOKEN,
-        IMG_CONTEXT_TOKEN,
-        VID_START_TOKEN,
-        VID_END_TOKEN,
-        VID_CONTEXT_TOKEN,
-        PATCH_START_TOKEN,
-        PATCH_END_TOKEN,
-        PATCH_CONTEXT_TOKEN,
-        AUD_START_TOKEN,
-        AUD_END_TOKEN,
-        AUD_CONTEXT_TOKEN,
-        QUAD_START_TOKEN,
-        QUAD_END_TOKEN,
-        REF_START_TOKEN,
-        REF_END_TOKEN,
-        BOX_START_TOKEN,
-        BOX_END_TOKEN,
-        IMG_TAG_TOKEN,
-        VID_TAG_TOKEN,
-        AUD_TAG_TOKEN,
-    ]
-    num_new_tokens = tokenizer.add_tokens(token_list, special_tokens=True)
-
-    token_list = [f"<|audio_{i}|>" for i in range(16384)]
-    num_new_tokens = tokenizer.add_tokens(token_list, special_tokens=False)
-
-    return tokenizer
-
-
-class SenseVoiceGLM4VoiceTokenizer:
-    def __init__(self, model_name_or_path, flow_path=None, rank=None):
-        self.model_name_or_path = model_name_or_path
-        self.flow_path = flow_path
-
-        if rank is None and torch.distributed.is_initialized():
-            rank = torch.distributed.get_rank()
-            self.rank = rank % 8
-        else:
-            self.rank = rank
-        logger.info(f"{self.rank=}")
-
-        self.sample_rate = 16000
-
-        self.is_discrete = True
-        self.is_contiguous = True
-
-        self._resample_buffer: dict[int, torchaudio.transforms.Resample] = {}
-        for sample_rate in [8000, 22050, 24000, 32000, 44100, 48000]:
-            print(f"torchaudio.transforms.Resample {sample_rate=} {self.sample_rate=}")
-            self._resample_buffer[sample_rate] = torchaudio.transforms.Resample(
-                orig_freq=sample_rate, new_freq=self.sample_rate
-            )
-
-    def load_model(self):
-
-        if hasattr(self, "whisper_model"):
-            return
-        if self.rank is not None:
-            self.device = f"cuda:{self.rank}"
-            torch.cuda.set_device(self.rank)
-        else:
-            self.device = "cpu"
-
-        logger.info(f"{self.device=} Loading SenseVoiceSmall")
-        from huggingface_hub import snapshot_download
-        model_dir = snapshot_download(repo_id="FunAudioLLM/SenseVoiceSmall")
-        
-        _, self.kwargs = SenseVoiceSmall.from_pretrained(model=model_dir, device=self.device)
-        logger.info(f"{self.device=} Loading SenseVoiceSmall Done")
-
-        logger.info(f"{self.device=} Loading GLM4VoiceTokenizer")
-        self.whisper_model = (
-            WhisperVQEncoder.from_pretrained(self.model_name_or_path).eval().to(self.device)
-        )
-        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(self.model_name_or_path)
-        print(hasattr(self, "whisper_model"),'whisper_model')
-
-        if self.flow_path is not None:
-            flow_config = os.path.join(self.flow_path, "config.yaml")
-            flow_checkpoint = os.path.join(self.flow_path, "flow.pt")
-            hift_checkpoint = os.path.join(self.flow_path, "hift.pt")
-
-            # Flow & Hift
-            self.audio_decoder = AudioDecoder(
-                config_path=flow_config,
-                flow_ckpt_path=flow_checkpoint,
-                hift_ckpt_path=hift_checkpoint,
-                device=self.device,
-            )
-        logger.info(f"{self.device=} Loading GLM4VoiceTokenizer Done")
-
-    def encode(self, audio_path, is_discrete=False, is_contiguous=True, **kwargs):
-
-        import faulthandler
-        faulthandler.enable()
-
-        if not hasattr(self, "whisper_model"):
-            self.load_model()
-
-        assert not (is_discrete and is_contiguous)
-        assert is_discrete or is_contiguous
-
-        if is_discrete:
-            audio_tokens = self.extract_speech_token(
-                self.whisper_model, self.feature_extractor, [audio_path], device=self.device
-            )[0]
-            return audio_tokens
-
-        if is_contiguous:
-
-            audio, sample_rate = torchaudio.load(audio_path)
-            audio = audio.mean(0)
-            if sample_rate != self.sample_rate:
-                if sample_rate not in self._resample_buffer:
-                    print(f"torchaudio.transforms.Resample {sample_rate=} {self.sample_rate=} {self.device=}", flush=True)
-                    self._resample_buffer[sample_rate] = torchaudio.transforms.Resample(
-                        orig_freq=sample_rate, new_freq=self.sample_rate
-                    ).to(self.device)
-                audio = audio.to(self.device)
-                self._resample_buffer[sample_rate].to(self.device)
-                audio = self._resample_buffer[sample_rate](audio[None, :])[0, :]
-                audio = audio.cpu()
-
-            frontend = self.kwargs["frontend"]
-
-            speech, speech_lengths = extract_fbank(audio, data_type="sound", frontend=frontend)
-
-            speech = speech[0]
-
-            return speech
-
-    def decode(self, audio_tokens, option_steps=10, **kwargs):
-        if not hasattr(self, "whisper_model"):
-            self.load_model()
-
-        this_uuid = str(uuid.uuid4())
-        this_uuid = "abc"
-
-        tts_token = torch.tensor(audio_tokens, device=self.device).unsqueeze(0)
-
-        flow_prompt_speech_token = torch.zeros(1, 0, dtype=torch.int64).to(self.device)
-        prompt_speech_feat = torch.zeros(1, 0, 80).to(self.device)
-
-        tts_speech, tts_mel = self.audio_decoder.token2wav(
-            tts_token,
-            uuid=this_uuid,
-            prompt_token=flow_prompt_speech_token.to(self.device),
-            prompt_feat=prompt_speech_feat.to(self.device),
-            finalize=True,
-            option_steps=option_steps,
-        )
-        tts_speechs = []
-        tts_speechs.append(tts_speech.squeeze())
-        tts_speech = torch.cat(tts_speechs, dim=-1).cpu()
-
-        return tts_speech
-
-    def apply_to_role(self, role, **kwargs):
-        is_discrete = kwargs.get("is_discrete", False)
-        if is_discrete and role in ["assistant", "gpt"]:
-            return True
-
-        is_contiguous = kwargs.get("is_contiguous", False)
-        if is_contiguous and role in ["user", "human"]:
-            return True
-
-        return False
-
-    def extract_speech_token(self, model, feature_extractor, utts, device="cuda"):
-        with torch.no_grad():
-            audios, indices = [], []
-            for idx, utt in enumerate(utts):
-                if isinstance(utt, tuple):
-                    audio, sample_rate = utt
-                else:
-                    audio, sample_rate = torchaudio.load(utt)
-                audio = audio.to(device)
-                if sample_rate != 16000:
-                    if sample_rate not in self._resample_buffer:
-                        print(f"torchaudio.transforms.Resample {sample_rate=} {self.sample_rate=} {device=}")
-                        self._resample_buffer[sample_rate] = torchaudio.transforms.Resample(
-                            orig_freq=sample_rate, new_freq=16000
-                        ).to(device)
-                    self._resample_buffer[sample_rate].to(device)
-                    audio = self._resample_buffer[sample_rate](audio)
-
-                audio = audio[0]
-                audio = audio.cpu().numpy()
-                time_step = 0
-                while time_step * 16000 < audio.shape[0]:
-                    audio_segment = audio[time_step * 16000 : (time_step + 30) * 16000]
-                    audios.append(audio_segment)
-                    indices.append(idx)
-                    time_step += 30
-            pooling_kernel_size = model.config.pooling_kernel_size or 1
-            stride = (
-                model.conv1.stride[0]
-                * model.conv2.stride[0]
-                * pooling_kernel_size
-                * feature_extractor.hop_length
-            )
-            all_speech_tokens = [[] for _ in range(len(utts))]
-            batch_size = 128
-            for start in range(0, len(audios), batch_size):
-                features = feature_extractor(
-                    audios[start : start + batch_size],
-                    sampling_rate=16000,
-                    return_attention_mask=True,
-                    return_tensors="pt",
-                    device=device,
-                    padding="longest",
-                    pad_to_multiple_of=stride,
-                )
-                features = features.to(device=device)
-                outputs = model(**features)
-                speech_tokens = outputs.quantized_token_ids
-                attention_mask = features.attention_mask[
-                    :, :: model.conv1.stride[0] * model.conv2.stride[0]
-                ]
-                attention_mask = attention_mask[:, :: model.config.pooling_kernel_size]
-                assert attention_mask.shape == speech_tokens.shape
-                for i in range(len(speech_tokens)):
-                    idx = indices[start + i]
-                    speech_token = speech_tokens[i][attention_mask[i].bool()].tolist()
-                    all_speech_tokens[idx].extend(speech_token)
-
-            return all_speech_tokens

s2i_0.wav

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9b8bd63844e968ffd84e083d34411ea6ccfcbdf101df3d8f0853cb557428be2f
-size 593358

third_party/GLM-4-Voice/.gitignore

@@ -1,4 +0,0 @@
-*venv
-*.DS_Store
-*.idea/
-test*

third_party/GLM-4-Voice/LICENSE

@@ -1,201 +0,0 @@
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright 2024 GLM-4-Voice Model Team @ Zhipu AI
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.

third_party/GLM-4-Voice/README.md

@@ -1,159 +0,0 @@
-# GLM-4-Voice
-<p align="center">
-📄<a href="https://arxiv.org/abs/2412.02612" target="_blank"> Report </a> • 🤗 <a href="https://huggingface.co/THUDM/glm-4-voice-9b" target="_blank">HF Repo</a> • 🤖 <a href="https://modelscope.cn/studios/ZhipuAI/GLM-4-Voice-Demo" target="_blank">Demo</a> • 🐦 <a href="https://twitter.com/thukeg" target="_blank">Twitter</a>
-</p>
-
-Read this in [English](./README_en.md)
-
-GLM-4-Voice 是智谱 AI 推出的端到端语音模型。GLM-4-Voice 能够直接理解和生成中英文语音，进行实时语音对话，并且能够遵循用户的指令要求改变语音的情感、语调、语速、方言等属性。
-
-## Model Architecture
-![Model Architecture](./resources/architecture.jpeg)
-
-GLM-4-Voice 由三个部分组成：
-* GLM-4-Voice-Tokenizer: 通过在 [Whisper](https://github.com/openai/whisper) 的 Encoder 部分增加 Vector Quantization 并在 ASR 数据上有监督训练，将连续的语音输入转化为离散的 token。每秒音频平均只需要用 12.5 个离散 token 表示。
-* GLM-4-Voice-Decoder: 基于 [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) 的 Flow Matching 模型结构训练的支持流式推理的语音解码器，将离散化的语音 token 转化为连续的语音输出。最少只需要 10 个语音 token 即可开始生成，降低端到端对话延迟。
-* GLM-4-Voice-9B: 在 [GLM-4-9B](https://github.com/THUDM/GLM-4) 的基础上进行语音模态的预训练和对齐，从而能够理解和生成离散化的语音 token。
-
-预训练方面，为了攻克模型在语音模态下的智商和合成表现力两个难关，我们将 Speech2Speech 任务解耦合为“根据用户音频做出文本回复”和“根据文本回复和用户语音合成回复语音”两个任务，并设计两种预训练目标，分别基于文本预训练数据和无监督音频数据合成语音-文本交错数据以适配这两种任务形式。GLM-4-Voice-9B 在 GLM-4-9B 的基座模型基础之上，经过了数百万小时音频和数千亿 token 的音频文本交错数据预训练，拥有很强的音频理解和建模能力。
-
-对齐方面，为了支持高质量的语音对话，我们设计了一套流式思考架构：根据用户语音，GLM-4-Voice 可以流式交替输出文本和语音两个模态的内容，其中语音模态以文本作为参照保证回复内容的高质量，并根据用户的语音指令要求做出相应的声音变化，在最大程度保留语言模型智商的情况下仍然具有端到端建模的能力，同时具备低延迟性，最低只需要输出 20 个 token 便可以合成语音。
-
-## Model List
-
-|         Model         |       Type       |                                                                     Download                                                                     |
-|:---------------------:|:----------------:|:------------------------------------------------------------------------------------------------------------------------------------------------:|
-| GLM-4-Voice-Tokenizer | Speech Tokenizer | [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-tokenizer) [🤖 ModelScope](https://modelscope.cn/models/ZhipuAI/glm-4-voice-tokenizer) |
-|    GLM-4-Voice-9B     |    Chat Model    |        [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-9b) [🤖 ModelScope](https://modelscope.cn/models/ZhipuAI/glm-4-voice-9b)        |
-|  GLM-4-Voice-Decoder  |  Speech Decoder  |   [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-decoder) [🤖 ModelScope](https://modelscope.cn/models/ZhipuAI/glm-4-voice-decoder)   |
-
-## Usage
-我们提供了可以直接启动的 Web Demo。用户可以输入语音或文本，模型会同时给出语音和文字回复。
-
-![](resources/web_demo.png)
-
-### Preparation
-
-首先下载仓库
-```shell
-git clone --recurse-submodules https://github.com/THUDM/GLM-4-Voice
-cd GLM-4-Voice
-```
-然后安装依赖。也可以使用我们提供的镜像 `zhipuai/glm-4-voice:0.1` 以跳过这一步。
-```shell
-pip install -r requirements.txt
-```
-由于 Decoder 模型不支持通过 `transformers` 初始化，因此 checkpoint 需要单独下载。
-
-```shell
-# git 模型下载，请确保已安装 git-lfs
-git lfs install
-git clone https://huggingface.co/THUDM/glm-4-voice-decoder
-```
-
-### Launch Web Demo
-
-1. 启动模型服务
-
-```shell
-python model_server.py --host localhost --model-path THUDM/glm-4-voice-9b --port 10000 --dtype bfloat16 --device cuda:0
-```
-
-如果你需要使用 Int4 精度启动，请运行
-
-```shell
-python model_server.py --host localhost --model-path THUDM/glm-4-voice-9b --port 10000 --dtype int4 --device cuda:0
-```
-
-此命令会自动下载 `glm-4-voice-9b`。如果网络条件不好，也手动下载之后通过 `--model-path` 指定本地的路径。
-
-2. 启动 web 服务
-
-```shell
-python web_demo.py --tokenizer-path  THUDM/glm-4-voice-tokenizer --model-path THUDM/glm-4-voice-9b --flow-path ./glm-4-voice-decoder
-```
-
-即可在 http://127.0.0.1:8888 访问 web demo。
-
-此命令会自动下载 `glm-4-voice-tokenizer` 和 `glm-4-voice-9b`。 请注意，`glm-4-voice-decoder` 需要手动下载。
-
-如果网络条件不好，可以手动下载这三个模型之后通过 `--tokenizer-path`, `--flow-path` 和 `--model-path` 指定本地的路径。
-
-### Known Issues
-
-* Gradio 的流式音频播放效果不稳定。在生成完成后点击对话框中的音频质量会更高。
-
-## Cases
-
-我们提供了 GLM-4-Voice 的部分对话案例，包括控制情绪、改变语速、生成方言等。
-
-* 用轻柔的声音引导我放松
-
-https://github.com/user-attachments/assets/4e3d9200-076d-4c28-a641-99df3af38eb0
-
-* 用激动的声音解说足球比赛
-
-https://github.com/user-attachments/assets/0163de2d-e876-4999-b1bc-bbfa364b799b
-
-* 用哀怨的声音讲一个鬼故事
-
-https://github.com/user-attachments/assets/a75b2087-d7bc-49fa-a0c5-e8c99935b39a
-
-* 用东北话介绍一下冬天有多冷
-
-https://github.com/user-attachments/assets/91ba54a1-8f5c-4cfe-8e87-16ed1ecf4037
-
-* 用重庆话念“吃葡萄不吐葡萄皮”
-
-https://github.com/user-attachments/assets/7eb72461-9e84-4d8e-9c58-1809cf6a8a9b
-
-* 用北京话念一句绕口令
-
-https://github.com/user-attachments/assets/a9bb223e-9c0a-440d-8537-0a7f16e31651
-
-  * 加快语速
-
-https://github.com/user-attachments/assets/c98a4604-366b-4304-917f-3c850a82fe9f
-
-  * 再快一点
-
-https://github.com/user-attachments/assets/d5ff0815-74f8-4738-b0f1-477cfc8dcc2d
-
-## Acknowledgements
-
-本项目的部分代码来自：
-* [CosyVoice](https://github.com/FunAudioLLM/CosyVoice)
-* [transformers](https://github.com/huggingface/transformers)
-* [GLM-4](https://github.com/THUDM/GLM-4)
-
-## 协议
-
-+ GLM-4 模型的权重的使用则需要遵循 [模型协议](https://huggingface.co/THUDM/glm-4-voice-9b/blob/main/LICENSE)。
-
-+ 本开源仓库的代码则遵循 [Apache 2.0](LICENSE) 协议。
-
-## 引用
-
-```
-@misc{zeng2024glm4,
-      title={GLM-4-Voice: Towards Intelligent and Human-Like End-to-End Spoken Chatbot}, 
-      author={Aohan Zeng and Zhengxiao Du and Mingdao Liu and Kedong Wang and Shengmin Jiang and Lei Zhao and Yuxiao Dong and Jie Tang},
-      year={2024},
-      eprint={2412.02612},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2412.02612}, 
-}
-```
-
-```
-@misc{zeng2024scaling,
-      title={Scaling Speech-Text Pre-training with Synthetic Interleaved Data}, 
-      author={Aohan Zeng and Zhengxiao Du and Mingdao Liu and Lei Zhang and Shengmin Jiang and Yuxiao Dong and Jie Tang},
-      year={2024},
-      eprint={2411.17607},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2411.17607}, 
-}
-```

third_party/GLM-4-Voice/README_en.md

@@ -1,148 +0,0 @@
-# GLM-4-Voice
-<p align="center">
-📄<a href="https://arxiv.org/abs/2412.02612" target="_blank"> Report </a> • 🤗 <a href="https://huggingface.co/THUDM/glm-4-voice-9b" target="_blank">HF Repo</a> • 🤖 <a href="https://modelscope.cn/studios/ZhipuAI/GLM-4-Voice-Demo" target="_blank">Demo</a> • 🐦 <a href="https://twitter.com/thukeg" target="_blank">Twitter</a>
-</p>
-
-GLM-4-Voice is an end-to-end voice model launched by Zhipu AI. GLM-4-Voice can directly understand and generate Chinese and English speech, engage in real-time voice conversations, and change attributes such as emotion, intonation, speech rate, and dialect based on user instructions.
-
-## Model Architecture
-
-![Model Architecture](./resources/architecture.jpeg)
-We provide the three components of GLM-4-Voice:
-* GLM-4-Voice-Tokenizer: Trained by adding vector quantization to the encoder part of [Whisper](https://github.com/openai/whisper), converting continuous speech input into discrete tokens. Each second of audio is converted into 12.5 discrete tokens.
-* GLM-4-Voice-9B: Pre-trained and aligned on speech modality based on [GLM-4-9B](https://github.com/THUDM/GLM-4), enabling understanding and generation of discretized speech.
-* GLM-4-Voice-Decoder: A speech decoder supporting streaming inference, retrained based on [CosyVoice](https://github.com/FunAudioLLM/CosyVoice), converting discrete speech tokens into continuous speech output. Generation can start with as few as 10 audio tokens, reducing conversation latency.
-
-## Model List
-
-|         Model         |       Type       |                               Download                               |
-|:---------------------:|:----------------:|:--------------------------------------------------------------------:|
-| GLM-4-Voice-Tokenizer | Speech Tokenizer | [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-tokenizer) |
-|    GLM-4-Voice-9B     |    Chat Model    |    [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-9b)     |
-|  GLM-4-Voice-Decoder  |  Speech Decoder  |  [🤗 Huggingface](https://huggingface.co/THUDM/glm-4-voice-decoder)  |
-
-## Usage
-We provide a Web Demo that can be launched directly. Users can input speech or text, and the model will respond with both speech and text.
-
-![](resources/web_demo.png)
-
-### Preparation
-
-First, download the repository
-```shell
-git clone --recurse-submodules https://github.com/THUDM/GLM-4-Voice
-cd GLM-4-Voice
-```
-Then, install the dependencies. You can also use our pre-built docker image `zhipuai/glm-4-voice:0.1` to skip the step.
-```shell
-pip install -r requirements.txt
-```
-Since the Decoder model does not support initialization via `transformers`, the checkpoint needs to be downloaded separately.
-
-```shell
-# Git model download, please ensure git-lfs is installed
-git clone https://huggingface.co/THUDM/glm-4-voice-decoder
-```
-
-### Launch Web Demo
-
-1. Start the model server
-
-```shell
-python model_server.py --host localhost --model-path THUDM/glm-4-voice-9b --port 10000 --dtype bfloat16 --device cuda:0
-```
-
-If you need to launch with Int4 precision, run
-
-```shell
-python model_server.py --host localhost --model-path THUDM/glm-4-voice-9b --port 10000 --dtype int4 --device cuda:0
-```
-
-This command will automatically download `glm-4-voice-9b`. If network conditions are poor, you can manually download it and specify the local path using `--model-path`.
-
-2. Start the web service
-
-```shell
-python web_demo.py --tokenizer-path  THUDM/glm-4-voice-tokenizer --model-path THUDM/glm-4-voice-9b --flow-path ./glm-4-voice-decoder
-```
-
-You can access the web demo at [http://127.0.0.1:8888](http://127.0.0.1:8888).
-This command will automatically download `glm-4-voice-tokenizer` and `glm-4-voice-9b`. Please note that `glm-4-voice-decoder` needs to be downloaded manually.
-If the network connection is poor, you can manually download these three models and specify the local paths using `--tokenizer-path`, `--flow-path`, and `--model-path`.
-
-### Known Issues
-* Gradio’s streaming audio playback can be unstable. The audio quality will be higher when clicking on the audio in the dialogue box after generation is complete.
-
-## Examples
-We provide some dialogue cases for GLM-4-Voice, including emotion control, speech rate alteration, dialect generation, etc. (The examples are in Chinese.)
-
-* Use a gentle voice to guide me to relax
-
-https://github.com/user-attachments/assets/4e3d9200-076d-4c28-a641-99df3af38eb0
-
-* Use an excited voice to commentate a football match
-
-https://github.com/user-attachments/assets/0163de2d-e876-4999-b1bc-bbfa364b799b
-
-* Tell a ghost story with a mournful voice
-
-https://github.com/user-attachments/assets/a75b2087-d7bc-49fa-a0c5-e8c99935b39a
-
-* Introduce how cold winter is with a Northeastern dialect
-
-https://github.com/user-attachments/assets/91ba54a1-8f5c-4cfe-8e87-16ed1ecf4037
-
-* Say "Eat grapes without spitting out the skins" in Chongqing dialect
-
-https://github.com/user-attachments/assets/7eb72461-9e84-4d8e-9c58-1809cf6a8a9b
-
-* Recite a tongue twister with a Beijing accent
-
-https://github.com/user-attachments/assets/a9bb223e-9c0a-440d-8537-0a7f16e31651
-
-  * Increase the speech rate
-
-https://github.com/user-attachments/assets/c98a4604-366b-4304-917f-3c850a82fe9f
-
-  * Even faster
-
-https://github.com/user-attachments/assets/d5ff0815-74f8-4738-b0f1-477cfc8dcc2d
-
-## Acknowledgements
-
-Some code in this project is from:
-* [CosyVoice](https://github.com/FunAudioLLM/CosyVoice)
-* [transformers](https://github.com/huggingface/transformers)
-* [GLM-4](https://github.com/THUDM/GLM-4)
-
-## License Agreement
-
-+ The use of GLM-4 model weights must follow the [Model License Agreement](https://huggingface.co/THUDM/glm-4-voice-9b/blob/main/LICENSE).
-
-+ The code in this open-source repository is licensed under the [Apache 2.0](LICENSE) License.
-
-## Citation
-
-```
-@misc{zeng2024glm4,
-      title={GLM-4-Voice: Towards Intelligent and Human-Like End-to-End Spoken Chatbot}, 
-      author={Aohan Zeng and Zhengxiao Du and Mingdao Liu and Kedong Wang and Shengmin Jiang and Lei Zhao and Yuxiao Dong and Jie Tang},
-      year={2024},
-      eprint={2412.02612},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2412.02612}, 
-}
-```
-
-```
-@misc{zeng2024scaling,
-      title={Scaling Speech-Text Pre-training with Synthetic Interleaved Data}, 
-      author={Aohan Zeng and Zhengxiao Du and Mingdao Liu and Lei Zhang and Shengmin Jiang and Yuxiao Dong and Jie Tang},
-      year={2024},
-      eprint={2411.17607},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2411.17607}, 
-}
-```

third_party/GLM-4-Voice/audio_process.py

@@ -1,93 +0,0 @@
-import os
-import librosa
-import soundfile as sf
-import numpy as np
-from pathlib import Path
-import io
-
-# Split audio stream at silence points to prevent playback stuttering issues
-# caused by AAC encoder frame padding when streaming audio through Gradio audio components.
-class AudioStreamProcessor:
-    def __init__(self, sr=22050, min_silence_duration=0.1, threshold_db=-40):
-        self.sr = sr
-        self.min_silence_duration = min_silence_duration
-        self.threshold_db = threshold_db
-        self.buffer = np.array([])
-  
-    
-    def process(self, audio_data, last=False):
-        """
-        Add audio data and process it
-        params:
-            audio_data: audio data in numpy array
-            last: whether this is the last chunk of data
-        returns:
-            Processed audio data, returns None if no split point is found
-        """
-
-        # Add new data to buffer
-        self.buffer = np.concatenate([self.buffer, audio_data]) if len(self.buffer) > 0 else audio_data
-        
-        if last:
-            result = self.buffer
-            self.buffer = np.array([])
-            return self._to_wav_bytes(result)
-            
-        # Find silence boundary
-        split_point = self._find_silence_boundary(self.buffer)
-        
-        if split_point is not None:
-            # Modified: Extend split point to the end of silence
-            silence_end = self._find_silence_end(split_point)
-            result = self.buffer[:silence_end]
-            self.buffer = self.buffer[silence_end:]
-            return self._to_wav_bytes(result)
-            
-        return None
-        
-    def _find_silence_boundary(self, audio):
-        """
-        Find the starting point of silence boundary in audio
-        """
-        # Convert audio to decibels
-        db = librosa.amplitude_to_db(np.abs(audio), ref=np.max)
-        
-        # Find points below threshold
-        silence_points = np.where(db < self.threshold_db)[0]
-        
-        if len(silence_points) == 0:
-            return None
-            
-        # Calculate minimum silence samples
-        min_silence_samples = int(self.min_silence_duration * self.sr)
-        
-        # Search backwards for continuous silence segment starting point
-        for i in range(len(silence_points) - min_silence_samples, -1, -1):
-            if i < 0:
-                break
-            if np.all(np.diff(silence_points[i:i+min_silence_samples]) == 1):
-                return silence_points[i]
-                
-        return None
-        
-    def _find_silence_end(self, start_point):
-        """
-        Find the end point of silence segment
-        """
-        db = librosa.amplitude_to_db(np.abs(self.buffer[start_point:]), ref=np.max)
-        silence_points = np.where(db >= self.threshold_db)[0]
-        
-        if len(silence_points) == 0:
-            return len(self.buffer)
-            
-        return start_point + silence_points[0]
-      
-    def _to_wav_bytes(self, audio_data):
-        """
-        trans_to_wav_bytes
-        """
-        wav_buffer = io.BytesIO()
-        sf.write(wav_buffer, audio_data, self.sr, format='WAV')
-        return wav_buffer.getvalue()
-      
-    

third_party/GLM-4-Voice/cosyvoice/bin/inference.py

@@ -1,114 +0,0 @@
-# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import print_function
-
-import argparse
-import logging
-logging.getLogger('matplotlib').setLevel(logging.WARNING)
-import os
-
-import torch
-from torch.utils.data import DataLoader
-import torchaudio
-from hyperpyyaml import load_hyperpyyaml
-from tqdm import tqdm
-from cosyvoice.cli.model import CosyVoiceModel
-
-from cosyvoice.dataset.dataset import Dataset
-
-def get_args():
-    parser = argparse.ArgumentParser(description='inference with your model')
-    parser.add_argument('--config', required=True, help='config file')
-    parser.add_argument('--prompt_data', required=True, help='prompt data file')
-    parser.add_argument('--prompt_utt2data', required=True, help='prompt data file')
-    parser.add_argument('--tts_text', required=True, help='tts input file')
-    parser.add_argument('--llm_model', required=True, help='llm model file')
-    parser.add_argument('--flow_model', required=True, help='flow model file')
-    parser.add_argument('--hifigan_model', required=True, help='hifigan model file')
-    parser.add_argument('--gpu',
-                        type=int,
-                        default=-1,
-                        help='gpu id for this rank, -1 for cpu')
-    parser.add_argument('--mode',
-                        default='sft',
-                        choices=['sft', 'zero_shot'],
-                        help='inference mode')
-    parser.add_argument('--result_dir', required=True, help='asr result file')
-    args = parser.parse_args()
-    print(args)
-    return args
-
-
-def main():
-    args = get_args()
-    logging.basicConfig(level=logging.DEBUG,
-                        format='%(asctime)s %(levelname)s %(message)s')
-    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
-
-    # Init cosyvoice models from configs
-    use_cuda = args.gpu >= 0 and torch.cuda.is_available()
-    device = torch.device('cuda' if use_cuda else 'cpu')
-    with open(args.config, 'r') as f:
-        configs = load_hyperpyyaml(f)
-
-    model = CosyVoiceModel(configs['llm'], configs['flow'], configs['hift'])
-    model.load(args.llm_model, args.flow_model, args.hifigan_model)
-
-    test_dataset = Dataset(args.prompt_data, data_pipeline=configs['data_pipeline'], mode='inference', shuffle=False, partition=False, tts_file=args.tts_text, prompt_utt2data=args.prompt_utt2data)
-    test_data_loader = DataLoader(test_dataset, batch_size=None, num_workers=0)
-
-    del configs
-    os.makedirs(args.result_dir, exist_ok=True)
-    fn = os.path.join(args.result_dir, 'wav.scp')
-    f = open(fn, 'w')
-    with torch.no_grad():
-        for batch_idx, batch in tqdm(enumerate(test_data_loader)):
-            utts = batch["utts"]
-            assert len(utts) == 1, "inference mode only support batchsize 1"
-            text = batch["text"]
-            text_token = batch["text_token"].to(device)
-            text_token_len = batch["text_token_len"].to(device)
-            tts_text = batch["tts_text"]
-            tts_index = batch["tts_index"]
-            tts_text_token = batch["tts_text_token"].to(device)
-            tts_text_token_len = batch["tts_text_token_len"].to(device)
-            speech_token = batch["speech_token"].to(device)
-            speech_token_len = batch["speech_token_len"].to(device)
-            speech_feat = batch["speech_feat"].to(device)
-            speech_feat_len = batch["speech_feat_len"].to(device)
-            utt_embedding = batch["utt_embedding"].to(device)
-            spk_embedding = batch["spk_embedding"].to(device)
-            if args.mode == 'sft':
-                model_input = {'text': tts_text_token, 'text_len': tts_text_token_len,
-                               'llm_embedding': spk_embedding, 'flow_embedding': spk_embedding}
-            else:
-                model_input = {'text': tts_text_token, 'text_len': tts_text_token_len,
-                               'prompt_text': text_token, 'prompt_text_len': text_token_len,
-                               'llm_prompt_speech_token': speech_token, 'llm_prompt_speech_token_len': speech_token_len,
-                               'flow_prompt_speech_token': speech_token, 'flow_prompt_speech_token_len': speech_token_len,
-                               'prompt_speech_feat': speech_feat, 'prompt_speech_feat_len': speech_feat_len,
-                               'llm_embedding': utt_embedding, 'flow_embedding': utt_embedding}
-            model_output = model.inference(**model_input)
-            tts_key = '{}_{}'.format(utts[0], tts_index[0])
-            tts_fn = os.path.join(args.result_dir, '{}.wav'.format(tts_key))
-            torchaudio.save(tts_fn, model_output['tts_speech'], sample_rate=22050)
-            f.write('{} {}\n'.format(tts_key, tts_fn))
-            f.flush()
-    f.close()
-    logging.info('Result wav.scp saved in {}'.format(fn))
-
-
-if __name__ == '__main__':
-    main()

third_party/GLM-4-Voice/cosyvoice/bin/train.py

@@ -1,140 +0,0 @@
-# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import print_function
-import argparse
-import datetime
-import logging
-logging.getLogger('matplotlib').setLevel(logging.WARNING)
-from copy import deepcopy
-import torch
-import torch.distributed as dist
-# import deepspeed
-import pdb
-from hyperpyyaml import load_hyperpyyaml
-
-from torch.distributed.elastic.multiprocessing.errors import record
-
-from cosyvoice.utils.executor import Executor
-from cosyvoice.utils.train_utils import (
-    init_distributed,
-    init_dataset_and_dataloader,
-    init_optimizer_and_scheduler,
-    init_summarywriter, save_model,
-    wrap_cuda_model, check_modify_and_save_config)
-
-
-def get_args():
-    parser = argparse.ArgumentParser(description='training your network')
-    parser.add_argument('--train_engine',
-                        default='torch_ddp',
-                        choices=['torch_ddp', 'deepspeed'],
-                        help='Engine for paralleled training')
-    parser.add_argument('--model', required=True, help='model which will be trained')
-    parser.add_argument('--config', required=True, help='config file')
-    parser.add_argument('--train_data', required=True, help='train data file')
-    parser.add_argument('--cv_data', required=True, help='cv data file')
-    parser.add_argument('--checkpoint', help='checkpoint model')
-    parser.add_argument('--model_dir', required=True, help='save model dir')
-    parser.add_argument('--tensorboard_dir',
-                        default='tensorboard',
-                        help='tensorboard log dir')
-    parser.add_argument('--ddp.dist_backend',
-                        dest='dist_backend',
-                        default='nccl',
-                        choices=['nccl', 'gloo'],
-                        help='distributed backend')
-    parser.add_argument('--num_workers',
-                        default=0,
-                        type=int,
-                        help='num of subprocess workers for reading')
-    parser.add_argument('--prefetch',
-                        default=100,
-                        type=int,
-                        help='prefetch number')
-    parser.add_argument('--pin_memory',
-                        action='store_true',
-                        default=False,
-                        help='Use pinned memory buffers used for reading')
-    parser.add_argument('--deepspeed.save_states',
-                        dest='save_states',
-                        default='model_only',
-                        choices=['model_only', 'model+optimizer'],
-                        help='save model/optimizer states')
-    parser.add_argument('--timeout',
-                        default=30,
-                        type=int,
-                        help='timeout (in seconds) of cosyvoice_join.')
-    # parser = deepspeed.add_config_arguments(parser)
-    args = parser.parse_args()
-    return args
-
-
-@record
-def main():
-    args = get_args()
-    logging.basicConfig(level=logging.DEBUG,
-                        format='%(asctime)s %(levelname)s %(message)s')
-
-    override_dict = {k: None for k in ['llm', 'flow', 'hift'] if k != args.model}
-    with open(args.config, 'r') as f:
-        configs = load_hyperpyyaml(f, overrides=override_dict)
-    configs['train_conf'].update(vars(args))
-
-    # Init env for ddp
-    init_distributed(args)
-
-    # Get dataset & dataloader
-    train_dataset, cv_dataset, train_data_loader, cv_data_loader = \
-        init_dataset_and_dataloader(args, configs)
-
-    # Do some sanity checks and save config to arsg.model_dir
-    configs = check_modify_and_save_config(args, configs)
-
-    # Tensorboard summary
-    writer = init_summarywriter(args)
-
-    # load checkpoint
-    model = configs[args.model]
-    if args.checkpoint is not None:
-        model.load_state_dict(torch.load(args.checkpoint, map_location='cpu'))
-
-    # Dispatch model from cpu to gpu
-    model = wrap_cuda_model(args, model)
-
-    # Get optimizer & scheduler
-    model, optimizer, scheduler = init_optimizer_and_scheduler(args, configs, model)
-    # pdb.set_trace()
-    # Save init checkpoints
-    info_dict = deepcopy(configs['train_conf'])
-    save_model(model, 'init', info_dict)
-
-    # Get executor
-    executor = Executor()
-
-    # Start training loop
-    for epoch in range(info_dict['max_epoch']):
-        executor.epoch = epoch
-        train_dataset.set_epoch(epoch)
-        dist.barrier()
-        # try:
-        #     dist.barrier()
-        # except RuntimeError as e:
-        #     logging.info('except RuntimeError as e: {}'.format(e))
-        group_join = dist.new_group(backend="gloo", timeout=datetime.timedelta(seconds=args.timeout))
-        executor.train_one_epoc(model, optimizer, scheduler, train_data_loader, cv_data_loader, writer, info_dict, group_join)
-        dist.destroy_process_group(group_join)
-
-if __name__ == '__main__':
-    main()