diff --git a/asr_0.wav b/asr_0.wav deleted file mode 100644 index 544e9373f0e217aded51d68d992cc92f67a9ae2f..0000000000000000000000000000000000000000 --- a/asr_0.wav +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8cc5488d43dadd87e8fa93f8c2a418b6771b7a9201c87ef39fff7b10e5650803 -size 1108218 diff --git a/assets/svqa_0.jpg b/assets/svqa_0.jpg deleted file mode 100644 index c781f8bac3f020cf5c804cafa3239a360ae2442e..0000000000000000000000000000000000000000 --- a/assets/svqa_0.jpg +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:885b9b03f88681745a5078087a2811e2d2fef8b492e33ffaaf70b5e290cd3dea -size 185149 diff --git a/assets/vqa_0.png b/assets/vqa_0.png deleted file mode 100644 index 44e4c7629d4721a6cb2581245b27f9d8e60cdf32..0000000000000000000000000000000000000000 --- a/assets/vqa_0.png +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:891d33b80bb3fb09e599411b81ffe0fa2d1e204e2371e8f777633bc018dbc5b9 -size 268052 diff --git a/chat_template.txt b/chat_template.txt deleted file mode 100644 index 62eaa1ebacbf9bf9ab2d51b4e7d23e0d3daa697a..0000000000000000000000000000000000000000 --- a/chat_template.txt +++ /dev/null @@ -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 XML tags:\\n\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n\\n\\nFor each function call, return a json object with function name and arguments within XML tags:\\n\\n{\\\"name\\\": , \\\"arguments\\\": }\\n<|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\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {{- tool_call.arguments | tojson }}\n {{- '}\\n' }}\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\\n' }}\n {{- message.content }}\n {{- '\\n' }}\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 diff --git a/magvit_tokenizer/__init__.py b/magvit_tokenizer/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/magvit_tokenizer/common_modules.py b/magvit_tokenizer/common_modules.py deleted file mode 100644 index 150e715846b5be69faf74734d04d6ce148d41a1c..0000000000000000000000000000000000000000 --- a/magvit_tokenizer/common_modules.py +++ /dev/null @@ -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 diff --git a/magvit_tokenizer/misc.py b/magvit_tokenizer/misc.py deleted file mode 100644 index a7b919632431afd89b6ba9e7a99e7ea68ff35a31..0000000000000000000000000000000000000000 --- a/magvit_tokenizer/misc.py +++ /dev/null @@ -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 diff --git a/magvit_tokenizer/modeling_magvitv2.py b/magvit_tokenizer/modeling_magvitv2.py deleted file mode 100644 index d76d4b4cf7a9b34d84b67348e9bd64d2ad2e583a..0000000000000000000000000000000000000000 --- a/magvit_tokenizer/modeling_magvitv2.py +++ /dev/null @@ -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() diff --git a/magvit_tokenizer/modeling_utils.py b/magvit_tokenizer/modeling_utils.py deleted file mode 100644 index 0315ed0df25e6652e2307879d85b70a035deb2c8..0000000000000000000000000000000000000000 --- a/magvit_tokenizer/modeling_utils.py +++ /dev/null @@ -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. - - # - - # ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes - # precedent. - - # - - # 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..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. - - - - 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. - - - - 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) diff --git a/omni_diffusion/__init__.py b/omni_diffusion/__init__.py deleted file mode 100644 index 97d5ae0bf1b16efccf06444da2039b898ff6749a..0000000000000000000000000000000000000000 --- a/omni_diffusion/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# omni_diffusion package - inference only, no eager imports diff --git a/omni_diffusion/constants.py b/omni_diffusion/constants.py deleted file mode 100644 index a9130aa2941ff34b066ba464dfb43368b02f7e81..0000000000000000000000000000000000000000 --- a/omni_diffusion/constants.py +++ /dev/null @@ -1,105 +0,0 @@ -import logging - -logger = logging.getLogger(__name__) - - -if False: - IMG_TAG_TOKEN = "" - IMG_CONTEXT_TOKEN = "" - IMG_START_TOKEN = "" - IMG_END_TOKEN = "" - - VID_TAG_TOKEN = "