Skip to content

Refactored vits to match standardized output collection interface #44154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Aatman09 wants to merge 1 commit into from
Closed

Refactored vits to match standardized output collection interface #44154

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 44 additions & 52 deletions src/transformers/models/vits/modeling_vits.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
from ...modeling_utils import PreTrainedModel
from ...utils import auto_docstring, logging, torch_compilable_check
from .configuration_vits import VitsConfig
from ..utils.capture_outputs import capture_outputs


logger = logging.get_logger(__name__)
Expand Down Expand Up @@ -1091,12 +1092,10 @@ def forward(
hidden_states: torch.FloatTensor,
padding_mask: torch.FloatTensor,
attention_mask: torch.Tensor | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
) -> tuple | BaseModelOutput:
all_hidden_states = () if output_hidden_states else None
all_self_attentions = () if output_attentions else None

) -> BaseModelOutput:
# all_hidden_states = () if output_hidden_states else None
# all_self_attentions = () if output_attentions else None

attention_mask = create_bidirectional_mask(
config=self.config,
Expand All @@ -1109,8 +1108,8 @@ def forward(
synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)

for encoder_layer in self.layers:
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
# if output_hidden_states:
# all_hidden_states = all_hidden_states + (hidden_states,)

# add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
dropout_probability = np.random.uniform(0, 1)
Expand All @@ -1122,28 +1121,23 @@ def forward(
hidden_states,
attention_mask=attention_mask,
padding_mask=padding_mask,
output_attentions=output_attentions,
# output_attentions=output_attentions,
)
hidden_states = layer_outputs[0]

if skip_the_layer:
layer_outputs = (None, None)

if output_attentions:
all_self_attentions = all_self_attentions + (layer_outputs[1],)
# if output_attentions:
# all_self_attentions = all_self_attentions + (layer_outputs[1],)

hidden_states = hidden_states * padding_mask

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

if not return_dict:
return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
# if not return_dict:
# return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)

return BaseModelOutput(
last_hidden_state=hidden_states,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
)


Expand All @@ -1164,36 +1158,35 @@ def forward(
input_ids: torch.Tensor,
padding_mask: torch.FloatTensor,
attention_mask: torch.Tensor | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = True,
# output_attentions: bool | None = None,
# output_hidden_states: bool | None = None,
# return_dict: bool | None = True,
) -> tuple[torch.Tensor] | VitsTextEncoderOutput:
hidden_states = self.embed_tokens(input_ids) * math.sqrt(self.config.hidden_size)

encoder_outputs = self.encoder(
hidden_states=hidden_states,
padding_mask=padding_mask,
attention_mask=attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
# output_attentions=output_attentions,
# output_hidden_states=output_hidden_states,
# return_dict=return_dict,
)

last_hidden_state = encoder_outputs[0] if not return_dict else encoder_outputs.last_hidden_state

last_hidden_state = encoder_outputs[0] #if not return_dict else encoder_outputs.last_hidden_state
stats = self.project(last_hidden_state.transpose(1, 2)).transpose(1, 2) * padding_mask
prior_means, prior_log_variances = torch.split(stats, self.config.flow_size, dim=2)

if not return_dict:
outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]
return outputs
# if not return_dict:
# outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]
# return outputs

return VitsTextEncoderOutput(
last_hidden_state=last_hidden_state,
prior_means=prior_means,
prior_log_variances=prior_log_variances,
hidden_states=encoder_outputs.hidden_states,
attentions=encoder_outputs.attentions,
# hidden_states=encoder_outputs.hidden_states,
# attentions=encoder_outputs.attentions,
)


Expand All @@ -1203,6 +1196,7 @@ class VitsPreTrainedModel(PreTrainedModel):
base_model_prefix = "vits"
main_input_name = "input_ids"
supports_gradient_checkpointing = True
_can_capture_outputs = {"attentions" : VitsAttention , "hidden_states" : VitsEncoderLayer}

@torch.no_grad()
def _init_weights(self, module: nn.Module):
Expand Down Expand Up @@ -1266,19 +1260,19 @@ def __init__(self, config: VitsConfig):

# Initialize weights and apply final processing
self.post_init()

@capture_outputs
@auto_docstring
def forward(
self,
input_ids: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
speaker_id: int | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
# output_attentions: bool | None = None,
# output_hidden_states: bool | None = None,
# return_dict: bool | None = None,
labels: torch.FloatTensor | None = None,
**kwargs,
) -> tuple[Any] | VitsModelOutput:
) -> VitsModelOutput:
r"""
speaker_id (`int`, *optional*):
Which speaker embedding to use. Only used for multispeaker models.
Expand All @@ -1305,11 +1299,11 @@ def forward(
torch.Size([1, 45824])
```
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
# output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
# output_hidden_states = (
# output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
# )
# return_dict = return_dict if return_dict is not None else self.config.use_return_dict

if labels is not None:
raise NotImplementedError("Training of VITS is not supported yet.")
Expand All @@ -1333,15 +1327,13 @@ def forward(
input_ids=input_ids,
padding_mask=input_padding_mask,
attention_mask=attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,

)
hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
hidden_states = text_encoder_output.last_hidden_state
hidden_states = hidden_states.transpose(1, 2)
input_padding_mask = input_padding_mask.transpose(1, 2)
prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
prior_means = text_encoder_output.prior_means
prior_log_variances = text_encoder_output.prior_log_variances

if self.config.use_stochastic_duration_prediction:
log_duration = self.duration_predictor(
Expand Down Expand Up @@ -1385,16 +1377,16 @@ def forward(
waveform = waveform.squeeze(1)
sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)

if not return_dict:
outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
return outputs
# if not return_dict:
# outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
# return outputs

return VitsModelOutput(
waveform=waveform,
sequence_lengths=sequence_lengths,
spectrogram=spectrogram,
hidden_states=text_encoder_output.hidden_states,
attentions=text_encoder_output.attentions,
# hidden_states=text_encoder_output.hidden_states,
# attentions=text_encoder_output.attentions,
)


Expand Down
Loading