diff --git a/indextts/accel/gpt2_accel.py b/indextts/accel/gpt2_accel.py index f96800e..b44a541 100644 --- a/indextts/accel/gpt2_accel.py +++ b/indextts/accel/gpt2_accel.py @@ -1,7 +1,9 @@ import torch import torch.nn as nn from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions -from transformers.models.gpt2.modeling_gpt2 import Conv1D, GPT2Block, GPT2Model +from transformers.cache_utils import Cache, DynamicCache +from transformers.models.gpt2.modeling_gpt2 import Conv1D, GPT2Block +from indextts.gpt.transformers_gpt2 import GPT2Model from .attention import Attention diff --git a/indextts/gpt/model.py b/indextts/gpt/model.py index 5f3b25c..c92917f 100644 --- a/indextts/gpt/model.py +++ b/indextts/gpt/model.py @@ -5,7 +5,8 @@ import torch.nn as nn import torch.nn.functional as F import transformers -from transformers import GPT2Config, LogitsProcessorList +from transformers import GPT2Config, LogitsProcessorList, DynamicCache +from transformers.cache_utils import Cache from indextts.gpt.transformers_gpt2 import GPT2PreTrainedModel, GPT2Model # from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList @@ -91,8 +92,13 @@ class GPT2InferenceModel(GPT2PreTrainedModel): token_type_ids = kwargs.get("token_type_ids", None) # usually None if not self.kv_cache: past_key_values = None + + # 将旧格式的元组转换为 DynamicCache（支持 v4.53.0+） + if past_key_values is not None and not isinstance(past_key_values, Cache): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + # only last token for inputs_ids if past is defined in kwargs - if past_key_values: + if past_key_values is not None: input_ids = input_ids[:, -1].unsqueeze(-1) if token_type_ids is not None: token_type_ids = token_type_ids[:, -1].unsqueeze(-1) @@ -104,7 +110,7 @@ class GPT2InferenceModel(GPT2PreTrainedModel): # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 0) - if past_key_values: + if past_key_values is not None: position_ids = position_ids[:, -1].unsqueeze(-1) else: position_ids = None diff --git a/indextts/gpt/model_v2.py b/indextts/gpt/model_v2.py index 3de3ac6..f3c6c16 100644 --- a/indextts/gpt/model_v2.py +++ b/indextts/gpt/model_v2.py @@ -5,7 +5,8 @@ import torch.nn as nn import torch.nn.functional as F import transformers -from transformers import GPT2Config, LogitsProcessorList +from transformers import GPT2Config, LogitsProcessorList, DynamicCache +from transformers.cache_utils import Cache from indextts.gpt.transformers_gpt2 import GPT2PreTrainedModel, GPT2Model # from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList @@ -91,8 +92,13 @@ class GPT2InferenceModel(GPT2PreTrainedModel): token_type_ids = kwargs.get("token_type_ids", None) # usually None if not self.kv_cache: past_key_values = None + + # 将旧格式的元组转换为 DynamicCache（支持 v4.53.0+） + if past_key_values is not None and not isinstance(past_key_values, Cache): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + # only last token for inputs_ids if past is defined in kwargs - if past_key_values: + if past_key_values is not None: input_ids = input_ids[:, -1].unsqueeze(-1) if token_type_ids is not None: token_type_ids = token_type_ids[:, -1].unsqueeze(-1) @@ -104,7 +110,7 @@ class GPT2InferenceModel(GPT2PreTrainedModel): # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 0) - if past_key_values: + if past_key_values is not None: position_ids = position_ids[:, -1].unsqueeze(-1) else: position_ids = None diff --git a/indextts/gpt/transformers_gpt2.py b/indextts/gpt/transformers_gpt2.py index ab7fa96..66cc72f 100755 --- a/indextts/gpt/transformers_gpt2.py +++ b/indextts/gpt/transformers_gpt2.py @@ -29,6 +29,7 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers.activations import ACT2FN import transformers +from transformers.cache_utils import Cache, DynamicCache from indextts.gpt.transformers_generation_utils import GenerationMixin from indextts.gpt.transformers_modeling_utils import PreTrainedModel @@ -1023,9 +1024,12 @@ class GPT2Model(GPT2PreTrainedModel): if token_type_ids is not None: token_type_ids = token_type_ids.view(-1, input_shape[-1]) + # 支持 Cache 类（v4.53.0+）和旧版元组格式 if past_key_values is None: past_length = 0 past_key_values = tuple([None] * len(self.h)) + elif isinstance(past_key_values, Cache): + past_length = past_key_values.get_seq_length() else: past_length = past_key_values[0][0].size(-2) if position_ids is None: @@ -1107,7 +1111,15 @@ class GPT2Model(GPT2PreTrainedModel): all_self_attentions = () if output_attentions else None all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None all_hidden_states = () if output_hidden_states else None - for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): + + # 处理 Cache 类的迭代 + if isinstance(past_key_values, Cache): + # Cache 类通过 __getitem__ 访问各层 + past_key_values_iter = [past_key_values[i] if i < len(self.h) else None for i in range(len(self.h))] + else: + past_key_values_iter = past_key_values + + for i, (block, layer_past) in enumerate(zip(self.h, past_key_values_iter)): # Model parallel if self.model_parallel: torch.cuda.set_device(hidden_states.device) @@ -1175,6 +1187,11 @@ class GPT2Model(GPT2PreTrainedModel): if v is not None ) + # 如果输入是 Cache 类，输出也应该是 Cache 类 + if use_cache and presents is not None: + if isinstance(past_key_values, Cache): + presents = DynamicCache.from_legacy_cache(presents) + return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=presents,