已合并
feat: add wan22 modeling #2920
chenpeizhe创建于 8月4日
feat: add wan22 modeling #2920
已合并
共 3 个文件变更+733-259
| @@ -0,0 +1,195 @@ | |||||||
| 1 | +"""Wan2.2 text encoder component (UMT5/T5/MT5 wrapper); structure built by ``build_text_encoder``, weights via the unified ``GenerativeBaseModel.setup_weights`` path.""" | ||||||
| 2 | + | ||||||
| 3 | +import logging | ||||||
| 4 | + | ||||||
| 5 | +import torch | ||||||
| 6 | +import torch.nn as nn | ||||||
| 7 | + | ||||||
| 8 | +from mindspeed_mm.fsdp.models.base_model import GenerativeBaseModel | ||||||
| 9 | +from mindspeed_mm.fsdp.utils.device import get_device_type | ||||||
| 10 | +from mindspeed_mm.fsdp.utils.dtype import get_dtype | ||||||
| 11 | + | ||||||
| 12 | +logger = logging.getLogger(__name__) | ||||||
| 13 | + | ||||||
| 14 | + | ||||||
| 15 | +class _SimpleTextEncoder(nn.Module, GenerativeBaseModel): | ||||||
| 16 | + """Minimal text encoder wrapper loading UMT5/T5/MT5 from transformers and providing encode().""" | ||||||
| 17 | + | ||||||
| 18 | + TRANSFORMERS_MAPPING = { | ||||||
| 19 | + "T5": "T5EncoderModel", | ||||||
| 20 | + "MT5": "MT5EncoderModel", | ||||||
| 21 | + "UMT5": "UMT5EncoderModel", | ||||||
| 22 | + } | ||||||
| 23 | + | ||||||
| 24 | + def __init__( | ||||||
| 25 | + self, | ||||||
| 26 | + model, | ||||||
| 27 | + use_attention_mask=True, | ||||||
| 28 | + output_key="last_hidden_state", | ||||||
| 29 | + hidden_state_skip_layer=None, | ||||||
| 30 | + ucg_rate=None, | ||||||
| 31 | + from_pretrained=None, | ||||||
| 32 | + load_format=None, | ||||||
| 33 | + ): | ||||||
| 34 | + super().__init__() | ||||||
| 35 | + self.model = model | ||||||
| 36 | + self.use_attention_mask = use_attention_mask | ||||||
| 37 | + self.output_key = output_key | ||||||
| 38 | + self.hidden_state_skip_layer = hidden_state_skip_layer | ||||||
| 39 | + self.ucg_rate = ucg_rate | ||||||
| 40 | + self._from_pretrained = from_pretrained | ||||||
| 41 | + self._load_format = load_format | ||||||
| 42 | + | ||||||
| 43 | + @classmethod | ||||||
| 44 | + def from_pretrained(cls, config): | ||||||
| 45 | + raise NotImplementedError( | ||||||
| 46 | + "_SimpleTextEncoder is constructed via build_text_encoder()." | ||||||
| 47 | + ) | ||||||
| 48 | + | ||||||
| 49 | + @classmethod | ||||||
| 50 | + def _from_config(cls, config): | ||||||
| 51 | + raise NotImplementedError( | ||||||
| 52 | + "_SimpleTextEncoder is constructed via build_text_encoder()." | ||||||
| 53 | + ) | ||||||
| 54 | + | ||||||
| 55 | + def get_weight_load_module(self): | ||||||
| 56 | + # The checkpoint holds the wrapped transformers model's keys, not the | ||||||
| 57 | + # wrapper's ``model.``-prefixed ones. | ||||||
| 58 | + return self.model | ||||||
| 59 | + | ||||||
| 60 | + def resolve_weight_source(self, config): | ||||||
| 61 | + from mindspeed_mm.fsdp.checkpoint.convert import resolve_checkpoint_source | ||||||
| 62 | + | ||||||
| 63 | + return resolve_checkpoint_source( | ||||||
| 64 | + config, | ||||||
| 65 | + fallback_from_pretrained=self._from_pretrained, | ||||||
| 66 | + fallback_format=self._load_format, | ||||||
| 67 | + ) | ||||||
| 68 | + | ||||||
| 69 | + def post_load(self, loaded: bool): | ||||||
| 70 | + # UMT5/T5 checkpoints store the token embedding as ``shared.weight`` and | ||||||
| 71 | + # expect it to be tied to ``encoder.embed_tokens.weight``. Re-bind it so | ||||||
| 72 | + # the embedding is not left randomly initialised. | ||||||
| 73 | + if loaded and hasattr(self.model, "encoder") and hasattr(self.model, "shared"): | ||||||
| 74 | + self.model.encoder.embed_tokens.weight = self.model.shared.weight | ||||||
| 75 | + logger.info( | ||||||
| 76 | + "[WAN-TE] %s ready on %s (weights_loaded=%s)", | ||||||
| 77 | + type(self.model).__name__, | ||||||
| 78 | + next(self.model.parameters()).device, | ||||||
| 79 | + loaded, | ||||||
| 80 | + ) | ||||||
| 81 | + | ||||||
| 82 | + def init_weights(self): | ||||||
| 83 | + # Guard for the trainer's meta-init stage: once pretrained weights are | ||||||
| 84 | + # loaded, re-initialization must be a no-op. Only reachable on the | ||||||
| 85 | + # random-init fallback path, where we delegate to the wrapped model. | ||||||
| 86 | + if getattr(self, "_weights_loaded", False): | ||||||
| 87 | + return | ||||||
| 88 | + if hasattr(self.model, "init_weights"): | ||||||
| 89 | + self.model.init_weights() | ||||||
| 90 | + | ||||||
| 91 | + def encode(self, input_ids, attention_mask, **kwargs): | ||||||
| 92 | + compute_device = get_device_type() | ||||||
| 93 | + # The wrapped T5 may be kept on CPU to save NPU memory. Move inputs to the | ||||||
| 94 | + # model's device for the forward pass and bring the outputs back to the | ||||||
| 95 | + # compute device so the rest of the pipeline stays on the NPU. | ||||||
| 96 | + first_param = next(self.model.parameters()) | ||||||
| 97 | + model_device = first_param.device | ||||||
| 98 | + *BN, L = input_ids.shape | ||||||
| 99 | + input_ids = input_ids.to(model_device).view(-1, L) | ||||||
| 100 | + attention_mask = attention_mask.to(model_device).view(-1, L) | ||||||
| 101 | + model_attention_mask = attention_mask if self.use_attention_mask else None | ||||||
| 102 | + | ||||||
| 103 | + output = self.model( | ||||||
| 104 | + input_ids=input_ids, | ||||||
| 105 | + attention_mask=model_attention_mask, | ||||||
| 106 | + output_hidden_states=self.hidden_state_skip_layer is not None, | ||||||
| 107 | + ) | ||||||
| 108 | + | ||||||
| 109 | + emb = output[self.output_key] | ||||||
| 110 | + if self.hidden_state_skip_layer is not None: | ||||||
| 111 | + emb = emb[-(self.hidden_state_skip_layer + 1)] | ||||||
 | |||||||
| 112 | + | ||||||
| 113 | + if self.ucg_rate is not None and self.ucg_rate > 0.0: | ||||||
🟡 Medium Priority 变更行：第 112、120 行 问题链： 触发条件：用户在配置中错误设置 建议：在 L112 的条件中增加上界校验 改动建议
  不准确？ | |||||||
| 114 | + def expand_dims_like(x, y): | ||||||
| 115 | + while x.dim() != y.dim(): | ||||||
| 116 | + x = x.unsqueeze(-1) | ||||||
| 117 | + return x | ||||||
| 118 | + emb = ( | ||||||
| 119 | + expand_dims_like( | ||||||
| 120 | + torch.bernoulli( | ||||||
| 121 | + (1.0 - self.ucg_rate) * torch.ones(emb.shape[0], device=emb.device, dtype=emb.dtype) | ||||||
| 122 | + ), | ||||||
| 123 | + emb, | ||||||
| 124 | + ) | ||||||
| 125 | + * emb | ||||||
| 126 | + ) | ||||||
| 127 | + | ||||||
| 128 | + if self.output_key in ["last_hidden_state", "hidden_states"]: | ||||||
| 129 | + emb = emb.view(*BN, emb.shape[-2], -1) | ||||||
| 130 | + elif self.output_key in ["pooler_output", "text_embeds"]: | ||||||
| 131 | + emb = emb.view(*BN, -1) | ||||||
| 132 | + else: | ||||||
| 133 | + raise NotImplementedError(f"Text encoder output_key: {self.output_key} is not implemented!") | ||||||
| 134 | + | ||||||
| 135 | + attention_mask = attention_mask.view(*BN, -1) | ||||||
| 136 | + return emb.to(compute_device), attention_mask.to(compute_device) | ||||||
| 137 | + | ||||||
| 138 | + def forward(self, input_ids, attention_mask, **kwargs): | ||||||
| 139 | + """Expose encode() as the module forward so FSDP2 hooks fire.""" | ||||||
| 140 | + return self.encode(input_ids, attention_mask, **kwargs) | ||||||
| 141 | + | ||||||
| 142 | + | ||||||
| 143 | +def build_text_encoder(config: dict): | ||||||
| 144 | + """Build a ``_SimpleTextEncoder`` structure on host CPU; weights are deferred to the unified ``GenerativeBaseModel.setup_weights`` path.""" | ||||||
| 145 | + from transformers import AutoConfig | ||||||
| 146 | + | ||||||
| 147 | + cfg = dict(config) # shallow copy | ||||||
| 148 | + backend = cfg.pop("hub_backend", "hf") | ||||||
| 149 | + if backend not in ("hf", "huggingface"): | ||||||
| 150 | + raise ValueError( | ||||||
| 151 | + f"Wan2.2 text_encoder only supports HuggingFace loading, got hub_backend={backend}." | ||||||
| 152 | + ) | ||||||
| 153 | + model_id = cfg.pop("model_id", "UMT5") | ||||||
| 154 | + if model_id not in _SimpleTextEncoder.TRANSFORMERS_MAPPING: | ||||||
| 155 | + raise ValueError(f"Model ID {model_id} is not supported for text encoder in pure FSDP2 mode") | ||||||
| 156 | + automodel_name = _SimpleTextEncoder.TRANSFORMERS_MAPPING[model_id] | ||||||
| 157 | + | ||||||
| 158 | + pretrained_path = cfg.pop("from_pretrained", None) | ||||||
| 159 | + if not pretrained_path: | ||||||
| 160 | + raise ValueError("Wan2.2 text_encoder requires 'from_pretrained' path.") | ||||||
| 161 | + torch_dtype = get_dtype(cfg.pop("dtype", "bf16")) | ||||||
| 162 | + | ||||||
| 163 | + logger.info("[WAN-TE] Loading %s config from %s", automodel_name, pretrained_path) | ||||||
| 164 | + transformer_config = AutoConfig.from_pretrained(pretrained_path) | ||||||
| 165 | + | ||||||
| 166 | + logger.info("[WAN-TE] Creating %s on CPU with dtype %s", automodel_name, torch_dtype) | ||||||
| 167 | + import transformers | ||||||
| 168 | + text_encoder = getattr(transformers, automodel_name)(transformer_config).to(torch_dtype) | ||||||
| 169 | + | ||||||
| 170 | + # Weight-source arbitration (load_format / checkpoint_format / hub_backend | ||||||
| 171 | + # aliases) is done by resolve_checkpoint_source at setup_weights time; the | ||||||
| 172 | + # raw values are only kept as fallbacks here. | ||||||
| 173 | + output_key = cfg.pop("output_key", "last_hidden_state") | ||||||
| 174 | + hidden_state_skip_layer = cfg.pop("hidden_state_skip_layer", None) | ||||||
| 175 | + ucg_rate = cfg.pop("ucg_rate", None) | ||||||
| 176 | + if hidden_state_skip_layer and output_key != "hidden_states": | ||||||
| 177 | + raise ValueError( | ||||||
| 178 | + f"hidden_state_skip_layer={hidden_state_skip_layer} requires output_key='hidden_states', " | ||||||
| 179 | + f"got {output_key!r}." | ||||||
| 180 | + ) | ||||||
| 181 | + if ucg_rate is not None and not 0.0 < ucg_rate <= 1.0: | ||||||
| 182 | + raise ValueError(f"ucg_rate must be in (0, 1], got {ucg_rate}.") | ||||||
| 183 | + wrapper = _SimpleTextEncoder( | ||||||
| 184 | + model=text_encoder, | ||||||
| 185 | + use_attention_mask=cfg.pop("use_attention_mask", True), | ||||||
| 186 | + output_key=output_key, | ||||||
| 187 | + hidden_state_skip_layer=hidden_state_skip_layer, | ||||||
| 188 | + ucg_rate=ucg_rate, | ||||||
| 189 | + from_pretrained=pretrained_path, | ||||||
| 190 | + load_format=cfg.pop("load_format", None) or cfg.pop("checkpoint_format", None), | ||||||
| 191 | + ) | ||||||
| 192 | + if isinstance(skip_to_device, str): | ||||||
| 193 | + skip_to_device = skip_to_device.strip().lower() in ("true", "1", "yes", "on") | ||||||
| 194 | + wrapper._ms_mm_skip_to_device = bool(skip_to_device) | ||||||
| 195 | + return wrapper | ||||||
| @@ -0,0 +1,267 @@ | |||
| 1 | +import importlib | ||
| 2 | +import warnings | ||
| 3 | +from abc import abstractmethod | ||
| 4 | + | ||
| 5 | +from einops import rearrange, repeat | ||
| 6 | +import torch | ||
| 7 | +import torch.nn as nn | ||
| 8 | +from diffusers.utils.accelerate_utils import apply_forward_hook | ||
| 9 | + | ||
| 10 | +from mindspeed_mm.fsdp.models.base_model import GenerativeBaseModel | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class DiffusersAEModel(nn.Module): | ||
| 14 | + """Diffusers autoencoder wrapper used by the FSDP2 Wan2.2 path.""" | ||
| 15 | + | ||
| 16 | + def __init__(self, model_name, config): | ||
| 17 | + super().__init__() | ||
| 18 | + module = importlib.import_module("diffusers") | ||
| 19 | + automodel = getattr(module, model_name) | ||
| 20 | + self.model_name = model_name | ||
| 21 | + self.model = automodel.from_pretrained(config["from_pretrained"], torch_dtype=config.pop("dtype", None)) | ||
| 22 | + self.do_sample = config.get("do_sample", True) | ||
| 23 | + | ||
| 24 | + self._tiling = False | ||
| 25 | + use_tiling = config.pop("enable_tiling", False) | ||
| 26 | + self.tiling_param = None | ||
| 27 | + if use_tiling: | ||
| 28 | + self.enable_tiling(tiling_param=config.get("tiling_param", None)) | ||
| 29 | + | ||
| 30 | + self.norm_latents = config.pop("norm_latents", True) | ||
| 31 | + self.norm_mode = config.pop("norm_mode", "value_shift_scale") | ||
| 32 | + | ||
| 33 | + def enable_tiling(self, tiling_param=None): | ||
| 34 | + has_custom_tiling_key = tiling_param is not None and ( | ||
| 35 | + "tile_size" in tiling_param or "tile_stride" in tiling_param | ||
| 36 | + ) | ||
| 37 | + if hasattr(self.model, "enable_tiling") and not has_custom_tiling_key: | ||
| 38 | + if tiling_param: | ||
| 39 | + self.model.enable_tiling(**tiling_param) | ||
| 40 | + self.tiling_param = tiling_param | ||
| 41 | + else: | ||
| 42 | + self.model.enable_tiling() | ||
| 43 | + else: | ||
| 44 | + self._tiling = True | ||
| 45 | + self.tiling_param = tiling_param if tiling_param else self.tiling_param | ||
| 46 | + warnings.warn( | ||
| 47 | + f"The autoencoder {self.model_name} in the diffusers library does not implement tiling functionality. " | ||
| 48 | + "Please ensure to call the custom tiling method to enable tiling. " | ||
| 49 | + ) | ||
| 50 | + | ||
| 51 | + def disable_tiling(self): | ||
| 52 | + if hasattr(self.model, "disable_tiling"): | ||
| 53 | + self.model.disable_tiling() | ||
| 54 | + else: | ||
| 55 | + self._tiling = False | ||
| 56 | + | ||
| 57 | + @apply_forward_hook | ||
| 58 | + def encode(self, x, **kwargs): | ||
| 59 | + if self._tiling: | ||
| 60 | + output = self.tiled_encode(x, **kwargs) | ||
| 61 | + else: | ||
| 62 | + output = self.model.encode(x, return_dict=True, **kwargs) | ||
| 63 | + if self.do_sample: | ||
| 64 | + output = output.latent_dist.sample() | ||
| 65 | + else: | ||
| 66 | + output = output.latent_dist.mode() | ||
| 67 | + | ||
| 68 | + if self.norm_latents: | ||
| 69 | + output = self.normalize_latent(output) | ||
| 70 | + | ||
| 71 | + return output | ||
| 72 | + | ||
| 73 | + # Alias so the ModelContainer can call the VAE as a normal module while | ||
| 74 | + # still using the decorated encode path. | ||
| 75 | + forward = encode | ||
| 76 | + | ||
| 77 | + @abstractmethod | ||
| 78 | + def tiled_encode(self, x, **kwargs): | ||
| 79 | + pass | ||
| 80 | + | ||
| 81 | + def normalize_latent(self, x): | ||
| 82 | + if self.norm_mode == "value_shift_scale": | ||
| 83 | + if getattr(self.model.config, "shift_factor", None): | ||
| 84 | + output = (x - self.model.config.shift_factor) * self.model.config.scale_factor | ||
| 85 | + else: | ||
| 86 | + output = x * self.model.config.scale_factor | ||
| 87 | + elif self.norm_mode == "channel_specified_shift_scale": | ||
| 88 | + latents_mean = torch.tensor(self.model.config.latents_mean).view(1, -1, 1, 1, 1).to(x) | ||
| 89 | + latents_std = torch.tensor(self.model.config.latents_std).view(1, -1, 1, 1, 1).to(x) | ||
| 90 | + output = (x - latents_mean) / latents_std | ||
| 91 | + else: | ||
| 92 | + raise NotImplementedError(f"norm_mode: {self.norm_mode} is not implemented.") | ||
| 93 | + return output | ||
| 94 | + | ||
| 95 | + @apply_forward_hook | ||
| 96 | + def decode(self, x, **kwargs): | ||
| 97 | + if self._tiling: | ||
| 98 | + return self.tiled_decode(x, **kwargs) | ||
| 99 | + return self.model.decode(x).sample | ||
| 100 | + | ||
| 101 | + @abstractmethod | ||
| 102 | + def tiled_decode(self, x, **kwargs): | ||
| 103 | + pass | ||
| 104 | + | ||
| 105 | + | ||
| 106 | +class WanVideoVAE(DiffusersAEModel, GenerativeBaseModel): | ||
| 107 | + """Wan2.2 VAE wrapper kept local to FSDP2 to avoid importing the Megatron AE package.""" | ||
| 108 | + | ||
| 109 | + def __init__(self, **config): | ||
| 110 | + super().__init__(model_name="AutoencoderKLWan", config=config) | ||
| 111 | + self.upsampling_factor = 8 | ||
| 112 | + | ||
| 113 | + @classmethod | ||
| 114 | + def from_pretrained(cls, config): | ||
| 115 | + return cls._from_config(config) | ||
| 116 | + | ||
| 117 | + @classmethod | ||
| 118 | + def _from_config(cls, config): | ||
| 119 | + cfg = config.to_dict() if hasattr(config, "to_dict") else dict(config) | ||
| 120 | + return cls(**cfg) | ||
| 121 | + | ||
| 122 | + def setup_weights(self, config): | ||
| 123 | + # Weights are loaded directly in ``__init__`` via diffusers | ||
| 124 | + # ``from_pretrained``; nothing to set up afterwards. The explicit | ||
| 125 | + # no-op prevents ModelContainerHub from re-loading the checkpoint | ||
| 126 | + # through the GenerativeBaseModel default (whose keys would not match this | ||
| 127 | + # wrapper's ``model.``-prefixed state dict). | ||
| 128 | + pass | ||
| 129 | + | ||
| 130 | + def _build_1d_mask(self, length, left_bound, right_bound, border_width): | ||
| 131 | + x = torch.ones((length,)) | ||
| 132 | + if not left_bound: | ||
| 133 | + x[:border_width] = (torch.arange(border_width) + 1) / border_width | ||
| 134 | + if not right_bound: | ||
| 135 | + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) | ||
| 136 | + return x | ||
| 137 | + | ||
| 138 | + def _build_mask(self, data, is_bound, border_width): | ||
| 139 | + _, _, _, h_size, w_size = data.shape | ||
| 140 | + h_mask = self._build_1d_mask(h_size, is_bound[0], is_bound[1], border_width[0]) | ||
| 141 | + w_mask = self._build_1d_mask(w_size, is_bound[2], is_bound[3], border_width[1]) | ||
| 142 | + | ||
| 143 | + h_mask = repeat(h_mask, "H -> H W", H=h_size, W=w_size) | ||
| 144 | + w_mask = repeat(w_mask, "W -> H W", H=h_size, W=w_size) | ||
| 145 | + | ||
| 146 | + mask = torch.stack([h_mask, w_mask]).min(dim=0).values | ||
| 147 | + mask = rearrange(mask, "H W -> 1 1 1 H W") | ||
| 148 | + return mask | ||
| 149 | + | ||
| 150 | + def tiled_encode(self, x, **kwargs): | ||
| 151 | + _, _, t_size, h_size, w_size = x.shape | ||
| 152 | + size_h, size_w = self.tiling_param["tile_size"] | ||
| 153 | + stride_h, stride_w = self.tiling_param["tile_stride"] | ||
| 154 | + size_h, size_w = ( | ||
| 155 | + size_h * self.upsampling_factor, | ||
| 156 | + size_w * self.upsampling_factor, | ||
| 157 | + ) | ||
| 158 | + stride_h, stride_w = ( | ||
| 159 | + stride_h * self.upsampling_factor, | ||
| 160 | + stride_w * self.upsampling_factor, | ||
| 161 | + ) | ||
| 162 | + | ||
| 163 | + tasks = [] | ||
| 164 | + for h_start in range(0, h_size, stride_h): | ||
| 165 | + if h_start - stride_h >= 0 and h_start - stride_h + size_h >= h_size: | ||
| 166 | + continue | ||
| 167 | + for w_start in range(0, w_size, stride_w): | ||
| 168 | + if w_start - stride_w >= 0 and w_start - stride_w + size_w >= w_size: | ||
| 169 | + continue | ||
| 170 | + h_end, w_end = h_start + size_h, w_start + size_w | ||
| 171 | + tasks.append((h_start, h_end, w_start, w_end)) | ||
| 172 | + | ||
| 173 | + out_t = (t_size + 3) // 4 | ||
| 174 | + weight = torch.zeros( | ||
| 175 | + (1, 1, out_t, h_size // self.upsampling_factor, w_size // self.upsampling_factor) | ||
| 176 | + ).to(x) | ||
| 177 | + values = torch.zeros( | ||
| 178 | + (1, 16, out_t, h_size // self.upsampling_factor, w_size // self.upsampling_factor) | ||
| 179 | + ).to(x) | ||
| 180 | + | ||
| 181 | + for h_start, h_end, w_start, w_end in tasks: | ||
| 182 | + hidden_states_batch = x[:, :, :, h_start:h_end, w_start:w_end] | ||
| 183 | + hidden_states_batch = self.model.encode(hidden_states_batch).latent_dist | ||
| 184 | + hidden_states_batch = hidden_states_batch.sample() if self.do_sample else hidden_states_batch.mode() | ||
| 185 | + | ||
| 186 | + mask = self._build_mask( | ||
| 187 | + hidden_states_batch, | ||
| 188 | + is_bound=(h_start == 0, h_end >= h_size, w_start == 0, w_end >= w_size), | ||
| 189 | + border_width=( | ||
| 190 | + (size_h - stride_h) // self.upsampling_factor, | ||
| 191 | + (size_w - stride_w) // self.upsampling_factor, | ||
| 192 | + ), | ||
| 193 | + ).to(x) | ||
| 194 | + | ||
| 195 | + target_h = h_start // self.upsampling_factor | ||
| 196 | + target_w = w_start // self.upsampling_factor | ||
| 197 | + values[ | ||
| 198 | + :, | ||
| 199 | + :, | ||
| 200 | + :, | ||
| 201 | + target_h:target_h + hidden_states_batch.shape[3], | ||
| 202 | + target_w:target_w + hidden_states_batch.shape[4], | ||
| 203 | + ] += hidden_states_batch * mask | ||
| 204 | + weight[ | ||
| 205 | + :, | ||
| 206 | + :, | ||
| 207 | + :, | ||
| 208 | + target_h:target_h + hidden_states_batch.shape[3], | ||
| 209 | + target_w:target_w + hidden_states_batch.shape[4], | ||
| 210 | + ] += mask | ||
| 211 | + values = values / weight | ||
| 212 | + return values | ||
| 213 | + | ||
| 214 | + def tiled_decode(self, x, **kwargs): | ||
| 215 | + _, _, t_size, h_size, w_size = x.shape | ||
| 216 | + size_h, size_w = self.tiling_param["tile_size"] | ||
| 217 | + stride_h, stride_w = self.tiling_param["tile_stride"] | ||
| 218 | + | ||
| 219 | + tasks = [] | ||
| 220 | + for h_start in range(0, h_size, stride_h): | ||
| 221 | + if h_start - stride_h >= 0 and h_start - stride_h + size_h >= h_size: | ||
| 222 | + continue | ||
| 223 | + for w_start in range(0, w_size, stride_w): | ||
| 224 | + if w_start - stride_w >= 0 and w_start - stride_w + size_w >= w_size: | ||
| 225 | + continue | ||
| 226 | + h_end, w_end = h_start + size_h, w_start + size_w | ||
| 227 | + tasks.append((h_start, h_end, w_start, w_end)) | ||
| 228 | + | ||
| 229 | + out_t = t_size * 4 - 3 | ||
| 230 | + weight = torch.zeros((1, 1, out_t, h_size * self.upsampling_factor, w_size * self.upsampling_factor)).to(x) | ||
| 231 | + values = torch.zeros((1, 3, out_t, h_size * self.upsampling_factor, w_size * self.upsampling_factor)).to(x) | ||
| 232 | + | ||
| 233 | + for h_start, h_end, w_start, w_end in tasks: | ||
| 234 | + hidden_states_batch = x[:, :, :, h_start:h_end, w_start:w_end] | ||
| 235 | + hidden_states_batch = self.model.decode(hidden_states_batch).sample | ||
| 236 | + | ||
| 237 | + mask = self._build_mask( | ||
| 238 | + hidden_states_batch, | ||
| 239 | + is_bound=(h_start == 0, h_end >= h_size, w_start == 0, w_end >= w_size), | ||
| 240 | + border_width=( | ||
| 241 | + (size_h - stride_h) * self.upsampling_factor, | ||
| 242 | + (size_w - stride_w) * self.upsampling_factor, | ||
| 243 | + ), | ||
| 244 | + ).to(x) | ||
| 245 | + | ||
| 246 | + target_h = h_start * self.upsampling_factor | ||
| 247 | + target_w = w_start * self.upsampling_factor | ||
| 248 | + values[ | ||
| 249 | + :, | ||
| 250 | + :, | ||
| 251 | + :, | ||
| 252 | + target_h:target_h + hidden_states_batch.shape[3], | ||
| 253 | + target_w:target_w + hidden_states_batch.shape[4], | ||
| 254 | + ] += hidden_states_batch * mask | ||
| 255 | + weight[ | ||
| 256 | + :, | ||
| 257 | + :, | ||
| 258 | + :, | ||
| 259 | + target_h:target_h + hidden_states_batch.shape[3], | ||
| 260 | + target_w:target_w + hidden_states_batch.shape[4], | ||
| 261 | + ] += mask | ||
| 262 | + values = values / weight | ||
| 263 | + values = values.clamp_(-1, 1) | ||
| 264 | + return values | ||
| 265 | + | ||
| 266 | + def get_tiling_state(self): | ||
| 267 | + return self._tiling | ||
🟡 Medium Priority
变更行：第 108-110 行
emb = output[self.output_key] if self.hidden_state_skip_layer: emb = emb[-(self.hidden_state_skip_layer + 1)]
问题链：
触发条件：配置
hidden_state_skip_layer > 0且output_key非"hidden_states"（包括默认的"last_hidden_state"）。建议：在 hidden_state_skip_layer 为真值时，从 output["hidden_states"] 取值进行层索引，而非从 output[self.output_key] 取值。修改方式：