mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-07-13 04:16:58 +08:00
Merge 4ece4eebfc6c79c857f7c67d71b3a306c458ae43 into 440268d3940eb14a01595439bbc05c4aacde9c72
This commit is contained in:
Yousef R. Gamaleldin
GitHub
commit
8ed3395adc
734
comfy/autoregressive_sampling.py
Normal file
734
comfy/autoregressive_sampling.py
Normal file
@ -0,0 +1,734 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import copy
|
||||
import torch
|
||||
import inspect
|
||||
import warnings
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, fields
|
||||
from typing import Optional, Union, Any
|
||||
from abc import ABC, abstractmethod
|
||||
import comfy.model_management
|
||||
|
||||
class CacheLayerParent(ABC):
|
||||
def __init__(self):
|
||||
self.keys, self.values = None, None
|
||||
|
||||
@abstractmethod
|
||||
def lazy_init(self, key_states): ...
|
||||
|
||||
@abstractmethod
|
||||
def update(self, key_states, value_states, cache_kwargs = None): ...
|
||||
|
||||
@abstractmethod
|
||||
def get_seq_length(self) -> int: ...
|
||||
|
||||
@abstractmethod
|
||||
def get_max_cache_shape(self) -> int: ...
|
||||
|
||||
def reset(self):
|
||||
if self.keys is not None:
|
||||
self.keys.zero_()
|
||||
self.values.zero_()
|
||||
|
||||
class Cache:
|
||||
def __init__(self):
|
||||
self.layers = []
|
||||
# TODO: offload logic
|
||||
def get_seq_length(self, layer_idx = 0):
|
||||
if layer_idx >= len(self.layers):
|
||||
return 0
|
||||
return self.layers[layer_idx].get_seq_length()
|
||||
|
||||
def get_max_cache_shape(self, layer_idx = 0):
|
||||
if layer_idx >= len(self.layers): # for future dynamic cache impl.
|
||||
return -1
|
||||
return self.layers[layer_idx].get_max_cache_shape()
|
||||
def reset(self):
|
||||
for layer_idx in range(len(self.layers)):
|
||||
self.layers[layer_idx].reset()
|
||||
|
||||
def update(self, key_states, value_states, layer_idx, cache_kwargs):
|
||||
keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
|
||||
return keys, values
|
||||
|
||||
@property
|
||||
def max_cache_len(self):
|
||||
values = [layer.max_cache_len for layer in self.layers]
|
||||
return max(values)
|
||||
|
||||
class StaticLayer(CacheLayerParent):
|
||||
def __init__(self, max_cache_len):
|
||||
super().__init__()
|
||||
self.max_cache_len = max_cache_len
|
||||
|
||||
def get_max_cache_shape(self):
|
||||
return self.max_cache_len
|
||||
|
||||
def get_seq_length(self):
|
||||
return (self.keys[0, 0].any(dim=-1)).sum() if self.keys is not None else 0
|
||||
|
||||
def lazy_init(self, key_states):
|
||||
|
||||
self.max_batch_size, self.num_heads, _, self.head_dim = key_states.shape
|
||||
self.dtype, self.device = key_states.dtype, key_states.device
|
||||
|
||||
self.keys = torch.zeros(
|
||||
(self.max_batch_size, self.num_heads, self.max_cache_len, self.head_dim),
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
self.values = torch.zeros(
|
||||
(self.max_batch_size, self.num_heads, self.max_cache_len, self.head_dim),
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
def update(self, key_states, value_states, cache_kwargs = None):
|
||||
|
||||
if self.keys is None:
|
||||
self.lazy_init(key_states)
|
||||
|
||||
cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None
|
||||
cache_position = (
|
||||
cache_position if cache_position is not None else torch.arange(key_states.shape[-2], device=self.device)
|
||||
)
|
||||
|
||||
try:
|
||||
self.keys.index_copy_(2, cache_position, key_states)
|
||||
self.values.index_copy_(2, cache_position, value_states)
|
||||
except NotImplementedError:
|
||||
# mps fallback
|
||||
self.keys[:, :, cache_position] = key_states
|
||||
self.values[:, :, cache_position] = value_states
|
||||
return self.keys, self.values
|
||||
|
||||
class StaticCache(Cache):
|
||||
def __init__(self, config, max_cache_len):
|
||||
self.layers = [StaticLayer(max_cache_len) for _ in range(config.num_hidden_layers)]
|
||||
|
||||
# TODO
|
||||
class DynamicCache(Cache):
|
||||
pass
|
||||
|
||||
NEED_SETUP_CACHE_CLASSES_MAPPING = { "static": StaticCache }
|
||||
|
||||
class TopKLogits:
|
||||
def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
|
||||
|
||||
self.top_k = max(top_k, min_tokens_to_keep)
|
||||
self.filter_value = filter_value
|
||||
|
||||
def __call__(self, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
top_k = min(self.top_k, scores.size(-1)) # cap top_k with vocab_size
|
||||
indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
|
||||
scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
|
||||
return scores_processed
|
||||
|
||||
class TemperatureLogitsWarper:
|
||||
def __init__(self, temperature: float):
|
||||
if not (temperature > 0):
|
||||
raise ValueError(f"`temperature` (={temperature}) must be a positive number > 0")
|
||||
|
||||
self.temperature = temperature
|
||||
|
||||
def __call__(self, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
scores_processed = scores / self.temperature
|
||||
return scores_processed
|
||||
|
||||
class TopPLogitsWarper:
|
||||
def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
|
||||
top_p = float(top_p)
|
||||
if top_p < 0 or top_p > 1.0:
|
||||
raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
|
||||
|
||||
self.top_p = top_p
|
||||
self.filter_value = filter_value
|
||||
self.min_tokens_to_keep = min_tokens_to_keep
|
||||
|
||||
def __call__(self, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
sorted_logits, sorted_indices = torch.sort(scores, descending=False)
|
||||
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
|
||||
|
||||
sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)
|
||||
sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0
|
||||
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
||||
scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
|
||||
return scores_processed
|
||||
|
||||
class MinLengthLogitsProcessor:
|
||||
def __init__(self, min_length: int, eos_token_id: torch.Tensor):
|
||||
self.min_length = min_length
|
||||
self.eos_token_id = eos_token_id
|
||||
|
||||
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
|
||||
if input_ids is None:
|
||||
return scores
|
||||
|
||||
vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
|
||||
eos_token_mask = torch.isin(vocab_tensor, self.eos_token_id)
|
||||
scores_processed = scores.clone()
|
||||
if input_ids.shape[-1] < self.min_length:
|
||||
scores_processed = torch.where(eos_token_mask, -math.inf, scores)
|
||||
return scores_processed
|
||||
|
||||
def get_logits_processing(config: GenerationConfig):
|
||||
# TODO: add support for beam search with diversity penalty
|
||||
logits_processors = []
|
||||
|
||||
if config._eos_token_tensor is not None and config.min_length > 1:
|
||||
logits_processors.append(MinLengthLogitsProcessor(config.min_length, config._eos_token_tensor))
|
||||
|
||||
if config.top_k is not None and config.top_k != 0:
|
||||
logits_processors.append(TopKLogits(config.top_k))
|
||||
|
||||
if config.temperature is not None and config.temperature != 1:
|
||||
logits_processors.append(TemperatureLogitsWarper(config.temperature))
|
||||
|
||||
if config.top_p is not None and config.top_p != 1:
|
||||
logits_processors.append(TopPLogitsWarper(config.top_p))
|
||||
|
||||
return logits_processors
|
||||
|
||||
def apply_logits_processing(input_ids, logits, logits_processing_list, **kwargs):
|
||||
for process in logits_processing_list:
|
||||
func_args = inspect.signature(process.__call__).parameters
|
||||
if not all(arg in kwargs for arg in list(func_args.keys())[3:]):
|
||||
raise ValueError(
|
||||
f"Make sure that all the required parameters: {list(func_args.keys())} for "
|
||||
f"{process.__class__} are passed to the logits processor."
|
||||
)
|
||||
if "input_ids" in func_args:
|
||||
logits = process(input_ids, logits)
|
||||
else:
|
||||
logits = process(logits, **kwargs)
|
||||
return logits
|
||||
|
||||
def check_stopping_strings(input_ids: torch.Tensor, stop_strings: list) -> torch.BoolTensor:
|
||||
# stop_strings must be a list of lists: List[List[], List[]]
|
||||
|
||||
device = input_ids.device
|
||||
batch_size, seq_len = input_ids.shape
|
||||
finished = torch.zeros(batch_size, dtype = torch.bool, device = device)
|
||||
|
||||
for b in range(batch_size):
|
||||
row = input_ids[b]
|
||||
# check each stop token sequence
|
||||
for stop_ids in stop_strings:
|
||||
n = len(stop_ids)
|
||||
if n == 0 or n > seq_len:
|
||||
continue
|
||||
# compare tail of the generated ids to the stop sequence
|
||||
if torch.all(row[-n:] == torch.tensor(stop_ids, device = device, dtype = row.dtype)):
|
||||
finished[b] = True
|
||||
break
|
||||
|
||||
return finished
|
||||
|
||||
def check_stopping_criteria(input_ids: torch.Tensor, max_length: int, eos_token, stop_strings: tuple = None):
|
||||
|
||||
device = input_ids.device
|
||||
|
||||
if not isinstance(eos_token, torch.Tensor):
|
||||
eos_token = torch.tensor(eos_token, device = device)
|
||||
|
||||
max_len_done = input_ids.shape[1] >= max_length
|
||||
|
||||
eos_done = torch.isin(input_ids[:, -1], eos_token)
|
||||
|
||||
if stop_strings is not None:
|
||||
stop_done = check_stopping_strings(input_ids, stop_strings)
|
||||
else:
|
||||
stop_done = torch.zeros(input_ids.size(0), dtype=torch.bool, device=device)
|
||||
|
||||
# finished either by lenght or eos or stop strings
|
||||
finished_mask = max_len_done | eos_done | stop_done
|
||||
|
||||
unfinished_mask = ~finished_mask
|
||||
|
||||
return finished_mask, unfinished_mask
|
||||
|
||||
class GenerationSampling(Enum):
|
||||
GREEDY_SEARCH = "greedy_search"
|
||||
BEAM_SAMPLING = "beam_sample"
|
||||
|
||||
@dataclass
|
||||
class GenerationConfig:
|
||||
pad_token_id: int = None
|
||||
bos_token_id: int = None
|
||||
eos_token_id: int = None
|
||||
|
||||
top_k: int = None
|
||||
top_p: int = None
|
||||
do_sample: bool = None # TODO: add beam sampling logic
|
||||
use_cache: bool = True
|
||||
num_beams: int = 1
|
||||
temperature: float = 1.0
|
||||
max_new_length: int = 1024
|
||||
min_new_length: int = 0
|
||||
num_return_sequences: int = 1
|
||||
generation_kwargs = {}
|
||||
cache_implementation: str = "static"
|
||||
is_using_cuda_graphs: bool = True
|
||||
|
||||
# special values, should be touched only in generation (no config)
|
||||
max_length = None
|
||||
min_length = None
|
||||
_bos_token_tensor = None
|
||||
_eos_token_tensor = None
|
||||
_pad_token_tensor = None
|
||||
|
||||
def update(self, **kwargs):
|
||||
to_remove = []
|
||||
for key, value in kwargs.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
to_remove.append(key)
|
||||
|
||||
# TODO in case of scaling for beam sampling: add a validation for the user's config
|
||||
|
||||
unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove}
|
||||
return unused_kwargs
|
||||
|
||||
@classmethod
|
||||
def from_model_config(cls, config_dict: dict, **kwargs) -> GenerationConfig:
|
||||
|
||||
config_dict = {key: value for key, value in config_dict.items() if value is not None}
|
||||
valid_fields = {f.name for f in fields(cls)}
|
||||
|
||||
filtered_args = {k: v for k, v in {**config_dict, **kwargs}.items() if k in valid_fields}
|
||||
|
||||
generation_config = cls(**filtered_args)
|
||||
|
||||
return generation_config
|
||||
|
||||
@dataclass
|
||||
class CacheConfig:
|
||||
hidden_size: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
head_dim: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
max_batch: int = 1
|
||||
|
||||
def get_text_config(self):
|
||||
return self
|
||||
|
||||
class AutoRegressiveGeneration:
|
||||
|
||||
def __init__(self, model, device, kv_cache_lengths: list = [1024, 4096, 8192]):
|
||||
|
||||
self.device = device
|
||||
self.dtype = model.dtype
|
||||
|
||||
self.model = model
|
||||
self.model.generation_config = GenerationConfig.from_model_config(self.model.config)
|
||||
self.model.generation_config.cache_implementation = self.model.cache_implementation
|
||||
|
||||
text_config = self.model.cache_config
|
||||
self.cache_config = CacheConfig(
|
||||
hidden_size = text_config["hidden_size"],
|
||||
num_attention_heads = text_config["num_attention_heads"],
|
||||
num_key_value_heads = text_config["num_key_value_heads"],
|
||||
head_dim = text_config["head_dim"],
|
||||
intermediate_size = text_config["intermediate_size"],
|
||||
num_hidden_layers = self.model.num_hidden_layers or text_config["num_hidden_layers"],
|
||||
)
|
||||
self.model.cache_config = self.cache_config
|
||||
|
||||
self.kv_caches = {
|
||||
length: StaticCache(
|
||||
config=self.cache_config,
|
||||
max_cache_len=length,
|
||||
)
|
||||
for length in sorted(kv_cache_lengths)
|
||||
|
||||
} if self.model.cache_implementation == "static" and self.model.use_kv_buckets else None
|
||||
|
||||
# for now
|
||||
self.model.generation_config.is_using_cuda_graphs = False
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids: Optional[torch.LongTensor] = None, max_new_length: int = 1024, min_new_length = 0,
|
||||
top_k: int = 50, top_p: float = 1.0, temperature: float = 1.0, do_sample: bool = False, seed = None, **kwargs):
|
||||
|
||||
if seed is not None:
|
||||
torch_generator = torch.Generator(device = input_ids.device).manual_seed(seed)
|
||||
else:
|
||||
torch_generator = None
|
||||
|
||||
kwargs["torch_generator"] = torch_generator
|
||||
kwargs["pbar"] = comfy.utils.ProgressBar(max_new_length)
|
||||
|
||||
gen_kwargs = dict(
|
||||
max_new_length = max_new_length,
|
||||
min_new_length = min_new_length,
|
||||
top_k = top_k,
|
||||
top_p = top_p,
|
||||
temperature = temperature,
|
||||
do_sample = do_sample
|
||||
)
|
||||
|
||||
# in case the specific model requires special handling
|
||||
generate_fn = getattr(self.model, "generate", None)
|
||||
if generate_fn is not None:
|
||||
generation_config = generate_fn(input_ids, generation_functions = self, **kwargs)
|
||||
generation_config.update(**gen_kwargs)
|
||||
|
||||
if generation_config is None:
|
||||
generation_config = GenerationConfig(max_new_length = max_new_length,
|
||||
min_new_length = min_new_length,
|
||||
top_k = top_k,
|
||||
top_p = top_p,
|
||||
do_sample = do_sample,
|
||||
temperature = temperature)
|
||||
|
||||
generation_config, model_kwargs = self._prepare_generation_config(
|
||||
generation_config, **kwargs
|
||||
)
|
||||
|
||||
accepts_attention_mask = "attention_mask" in set(inspect.signature(self.model.forward).parameters.keys())
|
||||
requires_attention_mask = "encoder_outputs" not in model_kwargs
|
||||
kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
|
||||
|
||||
inputs_tensor = input_ids
|
||||
model_input_name = "input_ids"
|
||||
batch_size = inputs_tensor.shape[0]
|
||||
|
||||
device = inputs_tensor.device
|
||||
self._prepare_special_tokens(generation_config, device = device)
|
||||
|
||||
if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask:
|
||||
model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
|
||||
inputs_tensor, generation_config, model_kwargs
|
||||
)
|
||||
elif kwargs_has_attention_mask:
|
||||
if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2:
|
||||
raise ValueError("`attention_mask` passed to `generate` must be 2D.")
|
||||
|
||||
input_ids_length = input_ids.shape[1]
|
||||
|
||||
generation_config = self._prepare_generated_length(
|
||||
generation_config=generation_config,
|
||||
input_ids_length=input_ids_length,
|
||||
)
|
||||
|
||||
self._validate_generated_length(generation_config, input_ids_length)
|
||||
|
||||
max_cache_length = generation_config.max_length - 1
|
||||
|
||||
self._prepare_cache_for_generation(
|
||||
generation_config, model_kwargs, batch_size, max_cache_length, device
|
||||
)
|
||||
|
||||
generation_mode = self.get_generation_mode(generation_config)
|
||||
|
||||
prepared_logits_processor = get_logits_processing(generation_config)
|
||||
|
||||
model_kwargs["use_cache"] = generation_config.use_cache
|
||||
|
||||
if generation_mode == GenerationSampling.GREEDY_SEARCH:
|
||||
input_ids, model_kwargs = self._expand_inputs_for_generation(
|
||||
input_ids = input_ids,
|
||||
expand_size = generation_config.num_return_sequences,
|
||||
**model_kwargs,
|
||||
)
|
||||
|
||||
# TODO: have a default self._sample fn and a default check if the model supports autoregGen or not
|
||||
if not hasattr(self.model, "_sample"):
|
||||
raise ValueError("Model doesn't support AutoRegressive Generation!")
|
||||
|
||||
self._prepare_kv_caches()
|
||||
|
||||
result = self.model._sample(
|
||||
input_ids,
|
||||
logits_processing_list = prepared_logits_processor,
|
||||
generation_config = generation_config,
|
||||
past_key_values_buckets = self.kv_caches,
|
||||
**model_kwargs,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def _prepare_kv_caches(self):
|
||||
for kv_cache in self.kv_caches.values():
|
||||
kv_cache.reset()
|
||||
|
||||
def get_generation_mode(self, config: GenerationConfig):
|
||||
if config.num_beams > 1:
|
||||
return GenerationSampling.BEAM_SAMPLING
|
||||
else:
|
||||
return GenerationSampling.GREEDY_SEARCH
|
||||
|
||||
def _prepare_generated_length(
|
||||
self,
|
||||
generation_config: GenerationConfig,
|
||||
input_ids_length,
|
||||
):
|
||||
|
||||
""" max_length = user_input_id_tokens + generation_max_length """
|
||||
|
||||
if generation_config.max_new_length is not None:
|
||||
generation_config.max_length = generation_config.max_new_length + input_ids_length
|
||||
|
||||
# same for max length
|
||||
if generation_config.min_new_length is not None:
|
||||
generation_config.min_length = generation_config.min_new_length + input_ids_length
|
||||
|
||||
return generation_config
|
||||
|
||||
def _get_cache(
|
||||
self, cache_implementation: str, batch_size: int, max_cache_len: int, device: torch.device, model_kwargs
|
||||
):
|
||||
|
||||
assert cache_implementation == "static", f"Only 'static' cache is supported, got {cache_implementation}"
|
||||
|
||||
cache_cls = NEED_SETUP_CACHE_CLASSES_MAPPING[cache_implementation]
|
||||
|
||||
if hasattr(self.model, "_cache"):
|
||||
cache_to_check = self.model._cache
|
||||
else:
|
||||
cache_to_check = None
|
||||
|
||||
need_new_cache = (
|
||||
not hasattr(self.model, "_cache")
|
||||
or (not isinstance(cache_to_check, cache_cls))
|
||||
or (self.cache_config.max_batch != batch_size)
|
||||
or (cache_to_check.max_cache_len < max_cache_len)
|
||||
)
|
||||
|
||||
if need_new_cache:
|
||||
cache_kwargs = {
|
||||
"config": self.cache_config,
|
||||
"max_cache_len": max_cache_len,
|
||||
}
|
||||
self.model._cache = cache_cls(**cache_kwargs)
|
||||
|
||||
else:
|
||||
self.model._cache.reset()
|
||||
|
||||
return self.model._cache
|
||||
|
||||
|
||||
def _prepare_cache_for_generation(
|
||||
self,
|
||||
generation_config: GenerationConfig,
|
||||
model_kwargs: dict,
|
||||
batch_size: int,
|
||||
max_cache_length: int,
|
||||
device: torch.device,
|
||||
) -> bool:
|
||||
|
||||
cache_name = "past_key_values"
|
||||
|
||||
if generation_config.use_cache is False:
|
||||
return
|
||||
|
||||
if generation_config.cache_implementation is not None:
|
||||
if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING:
|
||||
model_kwargs[cache_name] = self._get_cache(
|
||||
cache_implementation=generation_config.cache_implementation,
|
||||
batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
|
||||
max_cache_len=max_cache_length,
|
||||
device=device,
|
||||
model_kwargs=model_kwargs,
|
||||
)
|
||||
|
||||
elif generation_config.cache_implementation == "dynamic":
|
||||
model_kwargs[cache_name] = DynamicCache()
|
||||
|
||||
def _prepare_generation_config(self, generation_config: GenerationConfig, **kwargs: dict) -> tuple[GenerationConfig, dict]:
|
||||
|
||||
generation_config = copy.deepcopy(generation_config)
|
||||
|
||||
modified_values = {}
|
||||
global_default_generation_config = GenerationConfig()
|
||||
model_generation_config = self.model.generation_config
|
||||
|
||||
for key, model_gen_config_value in model_generation_config.__dict__.items():
|
||||
if key.startswith("_"):
|
||||
continue
|
||||
global_default_value = getattr(global_default_generation_config, key, None)
|
||||
custom_gen_config_value = getattr(generation_config, key, None)
|
||||
if (
|
||||
custom_gen_config_value == global_default_value
|
||||
and model_gen_config_value != global_default_value
|
||||
):
|
||||
modified_values[key] = model_gen_config_value
|
||||
setattr(generation_config, key, model_gen_config_value)
|
||||
|
||||
if generation_config.temperature == 0.0:
|
||||
generation_config.do_sample = False
|
||||
|
||||
model_kwargs = generation_config.update(**kwargs)
|
||||
|
||||
return generation_config, model_kwargs
|
||||
|
||||
def _validate_generated_length(self, generation_config: GenerationConfig, input_ids_length):
|
||||
"""Performs validation related to the resulting generated length"""
|
||||
|
||||
if input_ids_length >= generation_config.max_length:
|
||||
input_ids_string = "input_ids"
|
||||
raise ValueError(
|
||||
f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
|
||||
f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
|
||||
" increasing `max_length` or, better yet, setting `max_new_tokens`."
|
||||
)
|
||||
|
||||
min_length_error_suffix = (
|
||||
" Generation will stop at the defined maximum length. You should decrease the minimum length and/or "
|
||||
"increase the maximum length."
|
||||
)
|
||||
|
||||
if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
|
||||
warnings.warn(
|
||||
f"Unfeasible length constraints: `min_length` ({generation_config.min_length}) is larger than"
|
||||
f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix,
|
||||
UserWarning,
|
||||
)
|
||||
if generation_config.min_new_length is not None:
|
||||
min_length = generation_config.min_new_length + input_ids_length
|
||||
if min_length > generation_config.max_length:
|
||||
warnings.warn(
|
||||
f"Unfeasible length constraints: `min_new_length` ({generation_config.min_new_length}), when "
|
||||
f"added to the prompt length ({input_ids_length}), is larger than"
|
||||
f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix,
|
||||
UserWarning,
|
||||
)
|
||||
|
||||
def _expand_inputs_for_generation(
|
||||
self,
|
||||
expand_size: int = 1,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
**model_kwargs,
|
||||
) -> tuple[torch.LongTensor, dict[str, Any]]:
|
||||
""" Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...] """
|
||||
|
||||
if expand_size == 1:
|
||||
return input_ids, model_kwargs
|
||||
|
||||
def _expand_dict_for_generation(dict_to_expand):
|
||||
for key in dict_to_expand:
|
||||
if (
|
||||
key != "cache_position"
|
||||
and dict_to_expand[key] is not None
|
||||
and isinstance(dict_to_expand[key], torch.Tensor)
|
||||
):
|
||||
dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
|
||||
return dict_to_expand
|
||||
|
||||
if input_ids is not None:
|
||||
input_ids = input_ids.repeat_interleave(expand_size, dim=0)
|
||||
|
||||
model_kwargs = _expand_dict_for_generation(model_kwargs)
|
||||
|
||||
return input_ids, model_kwargs
|
||||
|
||||
def _prepare_special_tokens(
|
||||
self,
|
||||
generation_config: GenerationConfig,
|
||||
device: Optional[Union[torch.device, str]] = None,
|
||||
):
|
||||
|
||||
def _tensor_or_none(token, device=None):
|
||||
if token is None:
|
||||
return token
|
||||
|
||||
device = device if device is not None else self.device
|
||||
if isinstance(token, torch.Tensor):
|
||||
return token.to(device)
|
||||
return torch.tensor(token, device=device, dtype=torch.long)
|
||||
|
||||
bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
|
||||
eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
|
||||
pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
|
||||
|
||||
# We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
|
||||
if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
|
||||
eos_token_tensor = eos_token_tensor.unsqueeze(0)
|
||||
|
||||
if pad_token_tensor is None and eos_token_tensor is not None:
|
||||
pad_token_tensor = eos_token_tensor[0]
|
||||
|
||||
# to avoid problems with cuda graphs
|
||||
generation_config._bos_token_tensor = bos_token_tensor
|
||||
generation_config._eos_token_tensor = eos_token_tensor
|
||||
generation_config._pad_token_tensor = pad_token_tensor
|
||||
|
||||
def _prepare_attention_mask_for_generation(
|
||||
self,
|
||||
inputs_tensor: torch.Tensor,
|
||||
generation_config: GenerationConfig,
|
||||
model_kwargs: dict[str, Any],
|
||||
) -> torch.LongTensor:
|
||||
|
||||
pad_token_id = generation_config._pad_token_tensor
|
||||
eos_token_id = generation_config._eos_token_tensor
|
||||
|
||||
if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
|
||||
inputs_tensor = model_kwargs["input_ids"]
|
||||
|
||||
# No information for attention mask inference -> return default attention mask
|
||||
default_attention_mask = torch.ones(inputs_tensor.shape[:2], dtype=torch.long, device=inputs_tensor.device)
|
||||
if pad_token_id is None:
|
||||
return default_attention_mask
|
||||
|
||||
is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
|
||||
if not is_input_ids:
|
||||
return default_attention_mask
|
||||
|
||||
is_pad_token_in_inputs = (pad_token_id is not None) and (
|
||||
torch.isin(elements=inputs_tensor, test_elements=pad_token_id).any()
|
||||
)
|
||||
is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or ~(
|
||||
torch.isin(elements=eos_token_id, test_elements=pad_token_id).any()
|
||||
)
|
||||
can_infer_attention_mask = is_pad_token_in_inputs * is_pad_token_not_equal_to_eos_token_id
|
||||
attention_mask_from_padding = inputs_tensor.ne(pad_token_id).long()
|
||||
|
||||
attention_mask = (
|
||||
attention_mask_from_padding * can_infer_attention_mask + default_attention_mask * ~can_infer_attention_mask
|
||||
)
|
||||
return attention_mask
|
||||
|
||||
def auto_sample(node, patcher, input_ids, max_new_length=1024, min_new_length=0, top_k=50, top_p=1.0, temperature=1.0, do_sample = False, seed=None, **kwargs):
|
||||
# to work with BaseModel
|
||||
if hasattr(patcher, "model") and hasattr(patcher.model, "diffusion_model"):
|
||||
model = patcher.model.diffusion_model
|
||||
|
||||
if node._cached_autoregressive_sampler is None or node._cached_autoregressive_sampler.model is not model:
|
||||
if model.device != patcher.load_device:
|
||||
model = model.to(patcher.load_device, dtype=model.dtype)
|
||||
model.device = patcher.load_device
|
||||
node._cached_autoregressive_sampler = AutoRegressiveGeneration(model, device = patcher.load_device)
|
||||
|
||||
if isinstance(input_ids, dict):
|
||||
main_input_ids = input_ids.get("input_ids")
|
||||
kwargs.update({k: v for k, v in input_ids.items() if k != "input_ids"})
|
||||
else:
|
||||
main_input_ids = input_ids
|
||||
|
||||
device = node._cached_autoregressive_sampler.device
|
||||
|
||||
main_input_ids = main_input_ids.to(device)
|
||||
for k, v in kwargs.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
kwargs[k] = v.to(device)
|
||||
|
||||
# for tokenizers.Encoding output
|
||||
if not isinstance(main_input_ids, torch.Tensor):
|
||||
kwargs["attention_mask"] = torch.tensor(main_input_ids["attention_mask"])
|
||||
main_input_ids = torch.tensor(main_input_ids["input_ids"])
|
||||
|
||||
samples = node._cached_autoregressive_sampler.generate(main_input_ids, max_new_length, min_new_length, top_k, top_p, temperature, do_sample, seed=seed, **kwargs)
|
||||
if isinstance(samples, list):
|
||||
samples = [sample.to(comfy.model_management.intermediate_device()) for sample in samples]
|
||||
else:
|
||||
samples = samples.to(comfy.model_management.intermediate_device())
|
||||
return samples
|
||||
|
||||
311
comfy/ldm/higgsv2/loudness.py
Normal file
311
comfy/ldm/higgsv2/loudness.py
Normal file
@ -0,0 +1,311 @@
|
||||
import copy
|
||||
import math
|
||||
import torch
|
||||
import scipy
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from typing import Optional, List
|
||||
|
||||
# defaulted to the new pytorch api
|
||||
def _new_rfft(x: torch.Tensor):
|
||||
z = torch.fft.rfft(x, dim=-1)
|
||||
return torch.view_as_real(z)
|
||||
|
||||
def _new_irfft(x: torch.Tensor, length: int):
|
||||
x = torch.view_as_complex(x)
|
||||
return torch.fft.irfft(x, length, dim=-1)
|
||||
|
||||
def _compl_mul_conjugate(a: torch.Tensor, b: torch.Tensor):
|
||||
# changed this function to use the pytorch api
|
||||
return torch.view_as_real(torch.view_as_complex(a) * torch.view_as_complex(b).conj())
|
||||
|
||||
def unfold(input, kernel_size: int, stride: int):
|
||||
|
||||
shape = list(input.shape)
|
||||
length = shape.pop(-1)
|
||||
|
||||
n_frames = math.ceil((max(length, kernel_size) - kernel_size) / stride) + 1
|
||||
tgt_length = (n_frames - 1) * stride + kernel_size
|
||||
|
||||
padded = F.pad(input, (0, tgt_length - length)).contiguous()
|
||||
strides: List[int] = []
|
||||
|
||||
for dim in range(padded.dim()):
|
||||
strides.append(padded.stride(dim))
|
||||
|
||||
last_stride = strides.pop(-1)
|
||||
assert last_stride == 1, 'data should be contiguous'
|
||||
|
||||
strides = strides + [stride, 1]
|
||||
return padded.as_strided(shape + [n_frames, kernel_size], strides)
|
||||
|
||||
# convert the signal and filter to frequency domain, multiply them, then inverse FFT to get back to time-domain
|
||||
# faster than a sliding window over time-domain.
|
||||
def fft_conv1d(
|
||||
input: torch.Tensor, weight: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None, stride: int = 1, padding: int = 0,
|
||||
block_ratio: float = 5):
|
||||
|
||||
input = F.pad(input, (padding, padding))
|
||||
batch, _, length = input.shape
|
||||
out_channels, _, kernel_size = weight.shape
|
||||
|
||||
_rfft = _new_rfft
|
||||
_irfft = _new_irfft
|
||||
|
||||
if length < kernel_size:
|
||||
raise RuntimeError(f"Input should be at least as large as the kernel size {kernel_size}, "
|
||||
f"but it is only {length} samples long.")
|
||||
if block_ratio < 1:
|
||||
raise RuntimeError("Block ratio must be greater than 1.")
|
||||
|
||||
# We are going to process the input blocks by blocks, as for some reason it is faster
|
||||
# and less memory intensive (I think the culprit is `torch.einsum`.
|
||||
block_size: int = min(int(kernel_size * block_ratio), length)
|
||||
fold_stride = block_size - kernel_size + 1
|
||||
|
||||
# replaces to_pad
|
||||
weight = F.pad(weight, (0, block_size - weight.shape[-1]), mode = "constant", value = 0)
|
||||
weight_z = _rfft(weight)
|
||||
|
||||
# We pad the input and get the different frames, on which
|
||||
frames = unfold(input, block_size, fold_stride)
|
||||
|
||||
frames_z = _rfft(frames)
|
||||
out_z = _compl_mul_conjugate(frames_z, weight_z)
|
||||
out = _irfft(out_z, block_size)
|
||||
# The last bit is invalid, because FFT will do a circular convolution.
|
||||
out = out[..., :-kernel_size + 1]
|
||||
out = out.reshape(batch, out_channels, -1)
|
||||
out = out[..., ::stride]
|
||||
target_length = (length - kernel_size) // stride + 1
|
||||
out = out[..., :target_length]
|
||||
if bias is not None:
|
||||
out += bias[:, None]
|
||||
return out
|
||||
|
||||
class IIRfilter(object):
|
||||
|
||||
def __init__(self, G, Q, fc, rate, filter_type, passband_gain=1.0):
|
||||
G, Q, fc, rate = [t if isinstance(t, torch.Tensor) else torch.as_tensor(t) for t in (G, Q, fc, rate)]
|
||||
self.G = G
|
||||
self.Q = Q
|
||||
self.fc = fc
|
||||
self.rate = rate
|
||||
self.filter_type = filter_type
|
||||
self.passband_gain = passband_gain
|
||||
|
||||
def generate_coefficients(self):
|
||||
|
||||
A = 10**(self.G/40.0)
|
||||
w0 = 2.0 * torch.pi * (self.fc / self.rate)
|
||||
alpha = torch.sin(w0) / (2.0 * self.Q)
|
||||
|
||||
if self.filter_type == 'high_shelf':
|
||||
b0 = A * ( (A+1) + (A-1) * torch.cos(w0) + 2 * torch.sqrt(A) * alpha )
|
||||
b1 = -2 * A * ( (A-1) + (A+1) * torch.cos(w0) )
|
||||
b2 = A * ( (A+1) + (A-1) * torch.cos(w0) - 2 * torch.sqrt(A) * alpha )
|
||||
a0 = (A+1) - (A-1) * torch.cos(w0) + 2 * torch.sqrt(A) * alpha
|
||||
a1 = 2 * ( (A-1) - (A+1) * torch.cos(w0) )
|
||||
a2 = (A+1) - (A-1) * torch.cos(w0) - 2 * torch.sqrt(A) * alpha
|
||||
|
||||
elif self.filter_type == 'high_pass':
|
||||
b0 = (1 + torch.cos(w0))/2
|
||||
b1 = -(1 + torch.cos(w0))
|
||||
b2 = (1 + torch.cos(w0))/2
|
||||
a0 = 1 + alpha
|
||||
a1 = -2 * torch.cos(w0)
|
||||
a2 = 1 - alpha
|
||||
|
||||
return torch.tensor([b0, b1, b2])/a0, torch.tensor([a0, a1, a2])/a0
|
||||
|
||||
def apply_filter(self, data):
|
||||
return self.passband_gain * scipy.signal.lfilter(self.b, self.a, data)
|
||||
|
||||
@property
|
||||
def b_and_a(self):
|
||||
return self.generate_coefficients()
|
||||
|
||||
class Meter(torch.nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rate: int,
|
||||
filter_class: str = "K-weighting",
|
||||
block_size: float = 0.400,
|
||||
zeros: int = 512,
|
||||
use_fir: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.rate = rate
|
||||
self.filter_class = filter_class
|
||||
self.block_size = block_size
|
||||
self.use_fir = use_fir
|
||||
|
||||
G = torch.from_numpy(np.array([1.0, 1.0, 1.0, 1.41, 1.41]))
|
||||
self.register_buffer("G", G)
|
||||
|
||||
self._filters = {}
|
||||
self._filters['high_shelf'] = IIRfilter(4.0, 1/np.sqrt(2), 1500.0, self.rate, 'high_shelf')
|
||||
self._filters['high_pass'] = IIRfilter(0.0, 0.5, 38.0, self.rate, 'high_pass')
|
||||
|
||||
# Compute impulse responses so that filtering is fast via
|
||||
# a convolution at runtime, on GPU, unlike lfilter.
|
||||
impulse = np.zeros((zeros,))
|
||||
impulse[..., 0] = 1.0
|
||||
|
||||
firs = np.zeros((len(self._filters), 1, zeros))
|
||||
passband_gain = torch.tensor([filter.passband_gain for filter in self._filters.values()])
|
||||
|
||||
for i, (_, filter_stage) in enumerate(self._filters.items()):
|
||||
b, a = filter_stage.b_and_a
|
||||
firs[i] = scipy.signal.lfilter(b.numpy(), a.numpy(), impulse)
|
||||
|
||||
firs = torch.from_numpy(firs[..., ::-1].copy()).float()
|
||||
|
||||
self.register_buffer("firs", firs)
|
||||
self.register_buffer("passband_gain", passband_gain)
|
||||
|
||||
def apply_filter_fir(self, data: torch.Tensor):
|
||||
|
||||
# Data is of shape (nb, nch, nt)
|
||||
# Reshape to (nb*nch, 1, nt)
|
||||
nb, nt, nch = data.shape
|
||||
data = data.permute(0, 2, 1)
|
||||
data = data.reshape(nb * nch, 1, nt)
|
||||
|
||||
# Apply padding
|
||||
pad_length = self.firs.shape[-1]
|
||||
|
||||
# Apply filtering in sequence
|
||||
for i in range(self.firs.shape[0]):
|
||||
data = F.pad(data, (pad_length, pad_length))
|
||||
data = fft_conv1d(data, self.firs[i, None, ...])
|
||||
data = self.passband_gain[i] * data
|
||||
data = data[..., 1 : nt + 1]
|
||||
|
||||
data = data.permute(0, 2, 1)
|
||||
data = data[:, :nt, :]
|
||||
return data
|
||||
|
||||
def apply_filter(self, data: torch.Tensor):
|
||||
return self.apply_filter_fir(data)
|
||||
|
||||
def forward(self, data: torch.Tensor):
|
||||
return self.integrated_loudness(data)
|
||||
|
||||
def _unfold(self, input_data):
|
||||
T_g = self.block_size
|
||||
overlap = 0.75 # overlap of 75% of the block duration
|
||||
step = 1.0 - overlap # step size by percentage
|
||||
|
||||
kernel_size = int(T_g * self.rate)
|
||||
stride = int(T_g * self.rate * step)
|
||||
unfolded = unfold(input_data.permute(0, 2, 1), kernel_size, stride)
|
||||
unfolded = unfolded.transpose(-1, -2)
|
||||
|
||||
return unfolded
|
||||
|
||||
def integrated_loudness(self, data: torch.Tensor):
|
||||
|
||||
if not torch.is_tensor(data):
|
||||
data = torch.from_numpy(data).float()
|
||||
else:
|
||||
data = data.float()
|
||||
|
||||
input_data = copy.copy(data)
|
||||
# Data always has a batch and channel dimension.
|
||||
# Is of shape (nb, nt, nch)
|
||||
if input_data.ndim < 2:
|
||||
input_data = input_data.unsqueeze(-1)
|
||||
if input_data.ndim < 3:
|
||||
input_data = input_data.unsqueeze(0)
|
||||
|
||||
nb, _, nch = input_data.shape
|
||||
|
||||
# Apply frequency weighting filters - account
|
||||
# for the acoustic respose of the head and auditory system
|
||||
input_data = self.apply_filter(input_data)
|
||||
|
||||
G = self.G # channel gains
|
||||
T_g = self.block_size # 400 ms gating block standard
|
||||
Gamma_a = -70.0 # -70 LKFS = absolute loudness threshold
|
||||
|
||||
unfolded = self._unfold(input_data)
|
||||
|
||||
z = (1.0 / (T_g * self.rate)) * unfolded.square().sum(2)
|
||||
l = -0.691 + 10.0 * torch.log10((G[None, :nch, None] * z).sum(1, keepdim=True))
|
||||
l = l.expand_as(z)
|
||||
|
||||
# find gating block indices above absolute threshold
|
||||
z_avg_gated = z
|
||||
z_avg_gated[l <= Gamma_a] = 0
|
||||
masked = l > Gamma_a
|
||||
z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
|
||||
|
||||
# calculate the relative threshold value (see eq. 6)
|
||||
Gamma_r = (
|
||||
-0.691 + 10.0 * torch.log10((z_avg_gated * G[None, :nch]).sum(-1)) - 10.0
|
||||
)
|
||||
Gamma_r = Gamma_r[:, None, None]
|
||||
Gamma_r = Gamma_r.expand(nb, nch, l.shape[-1])
|
||||
|
||||
# find gating block indices above relative and absolute thresholds (end of eq. 7)
|
||||
z_avg_gated = z
|
||||
z_avg_gated[l <= Gamma_a] = 0
|
||||
z_avg_gated[l <= Gamma_r] = 0
|
||||
masked = (l > Gamma_a) * (l > Gamma_r)
|
||||
z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
|
||||
|
||||
# # Cannot use nan_to_num (pytorch 1.8 does not come with GCP-supported cuda version)
|
||||
# z_avg_gated = torch.nan_to_num(z_avg_gated)
|
||||
z_avg_gated = torch.where(
|
||||
z_avg_gated.isnan(), torch.zeros_like(z_avg_gated), z_avg_gated
|
||||
)
|
||||
z_avg_gated[z_avg_gated == float("inf")] = float(np.finfo(np.float32).max)
|
||||
z_avg_gated[z_avg_gated == -float("inf")] = float(np.finfo(np.float32).min)
|
||||
|
||||
LUFS = -0.691 + 10.0 * torch.log10((G[None, :nch] * z_avg_gated).sum(1))
|
||||
return LUFS.float()
|
||||
|
||||
|
||||
def loudness(
|
||||
audio_data, sample_rate: int, target_loudness: int, filter_class: str = "K-weighting", block_size: float = 0.400, **kwargs
|
||||
):
|
||||
MIN_LOUDNESS = -70
|
||||
device = audio_data.device
|
||||
|
||||
original_length = audio_data.shape[-1]
|
||||
signal_duration = original_length / sample_rate
|
||||
|
||||
# Pad if too short
|
||||
if signal_duration < 0.5:
|
||||
pad_len = int((0.5 - signal_duration) * sample_rate)
|
||||
audio_data = torch.nn.functional.pad(audio_data, (0, pad_len), mode="constant", value=0)
|
||||
|
||||
# create BS.1770 meter
|
||||
meter = Meter(
|
||||
sample_rate, filter_class=filter_class, block_size=block_size, **kwargs
|
||||
)
|
||||
meter = meter.to(audio_data.device)
|
||||
# measure loudness
|
||||
loudness = meter.integrated_loudness(audio_data.permute(0, 2, 1))
|
||||
audio_data = audio_data[..., :original_length]
|
||||
min_loudness = (
|
||||
torch.ones_like(loudness, device=loudness.device) * MIN_LOUDNESS
|
||||
)
|
||||
_loudness = torch.maximum(loudness, min_loudness)
|
||||
|
||||
_loudness = _loudness.to(device)
|
||||
|
||||
delta_loudness = target_loudness - _loudness
|
||||
gain = torch.pow(torch.tensor(10.0, device=device, dtype=audio_data.dtype), delta_loudness / 20.0)
|
||||
|
||||
output = gain * audio_data
|
||||
|
||||
if torch.max(torch.abs(output)) >= 1.0:
|
||||
import warnings
|
||||
warnings.warn("Possible clipped samples in output.")
|
||||
|
||||
return output
|
||||
1252
comfy/ldm/higgsv2/model.py
Normal file
1252
comfy/ldm/higgsv2/model.py
Normal file
File diff suppressed because it is too large
Load Diff
727
comfy/ldm/higgsv2/preprocess.py
Normal file
727
comfy/ldm/higgsv2/preprocess.py
Normal file
@ -0,0 +1,727 @@
|
||||
from typing import (
|
||||
Dict, List, Optional, Union, Tuple, MutableMapping, Any, Mapping, Collection, get_type_hints, get_args, get_origin
|
||||
)
|
||||
|
||||
from dataclasses import dataclass, fields, _FIELDS, _FIELD, _FIELD_INITVAR, is_dataclass
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import torch
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
__MAX_SIZE = 2048
|
||||
|
||||
def _ceil_to_nearest(n, round_to):
|
||||
return (n + round_to - 1) // round_to * round_to
|
||||
|
||||
@torch.jit.script
|
||||
def build_delay_pattern_mask(
|
||||
input_ids: torch.Tensor,
|
||||
bos_token_id: int,
|
||||
pad_token_id: int,
|
||||
):
|
||||
bsz, num_codebooks, seq_len = input_ids.shape
|
||||
|
||||
new_seq_len = seq_len + num_codebooks - 1
|
||||
input_ids_with_gen_mask = torch.ones((bsz, num_codebooks, new_seq_len), dtype=torch.long, device=input_ids.device)
|
||||
bos_mask = torch.tril(input_ids_with_gen_mask, -1) > 0
|
||||
eos_mask = torch.triu(input_ids_with_gen_mask, seq_len) > 0
|
||||
input_ids_with_gen_mask[bos_mask] = bos_token_id
|
||||
input_ids_with_gen_mask[(~bos_mask) & (~eos_mask)] = input_ids.reshape(-1)
|
||||
input_ids = input_ids_with_gen_mask.clone()
|
||||
input_ids[eos_mask] = pad_token_id
|
||||
input_ids_with_gen_mask[eos_mask] = -1
|
||||
return input_ids, input_ids_with_gen_mask
|
||||
|
||||
|
||||
# implementation of dacite's from_dict with only necessary parts for ChatML
|
||||
@lru_cache(maxsize=None)
|
||||
def cache(function):
|
||||
return lru_cache(maxsize=__MAX_SIZE, typed=True)(function)
|
||||
|
||||
@cache
|
||||
def get_fields(data_class):
|
||||
fields = getattr(data_class, _FIELDS)
|
||||
return [f for f in fields.values() if f._field_type is _FIELD or f._field_type is _FIELD_INITVAR]
|
||||
|
||||
def is_optional(type_) -> bool:
|
||||
return get_origin(type_) is Union and type(None) in get_args(type_)
|
||||
|
||||
def orig(data_class) -> Any:
|
||||
if is_dataclass(data_class):
|
||||
return data_class
|
||||
return get_origin(data_class)
|
||||
|
||||
@cache
|
||||
def extract_generic(type_, defaults: Tuple = ()) -> tuple:
|
||||
try:
|
||||
if getattr(type_, "_special", False):
|
||||
return defaults
|
||||
if type_.__args__ == ():
|
||||
return (type_.__args__,)
|
||||
return type_.__args__ or defaults # type: ignore
|
||||
except AttributeError:
|
||||
return defaults
|
||||
|
||||
def _build_value_for_collection(collection, data: Any) -> Any:
|
||||
if isinstance(data, Mapping):
|
||||
value_type = extract_generic(collection, defaults=(Any, Any))[1]
|
||||
return {
|
||||
key: _build_value(type_=value_type, data=value)
|
||||
for key, value in data.items()
|
||||
}
|
||||
|
||||
elif isinstance(data, Collection) and not isinstance(data, (str, bytes, Mapping)):
|
||||
item_type = extract_generic(collection, defaults=(Any,))[0]
|
||||
return [
|
||||
_build_value(type_=item_type, data=item)
|
||||
for item in data
|
||||
]
|
||||
|
||||
return data
|
||||
|
||||
def _build_value(type_, data) -> Any:
|
||||
if is_optional(type_) and data is None:
|
||||
return data
|
||||
if get_origin(type_) is Union:
|
||||
data = _build_value_for_union(union=type_, data=data)
|
||||
elif hasattr(type_, "__origin__"):
|
||||
data = _build_value_for_collection(collection=type_, data=data)
|
||||
elif cache(is_dataclass)(orig(type_)) and isinstance(data, Mapping):
|
||||
data = from_dict(data_class=type_, data=data)
|
||||
return data
|
||||
|
||||
def _build_value_for_union(union: type, data: Any) -> Any:
|
||||
for inner_type in get_args(union):
|
||||
if data is None and inner_type is type(None):
|
||||
return None
|
||||
try:
|
||||
return _build_value(inner_type, data)
|
||||
except Exception:
|
||||
continue
|
||||
raise ValueError(f"Cannot match {data!r} to any type in {union}")
|
||||
|
||||
def is_instance(value: Any, type_) -> bool:
|
||||
if type_ is Any:
|
||||
return True
|
||||
|
||||
origin = get_origin(type_)
|
||||
args = get_args(type_)
|
||||
|
||||
if origin is Union:
|
||||
return any(is_instance(value, arg) for arg in args)
|
||||
|
||||
if origin in (list, List):
|
||||
if not isinstance(value, list):
|
||||
return False
|
||||
(elem_type,) = args or (Any,)
|
||||
return all(is_instance(item, elem_type) for item in value)
|
||||
|
||||
if origin in (dict, Dict, Mapping):
|
||||
if not isinstance(value, dict):
|
||||
return False
|
||||
key_type, val_type = args or (Any, Any)
|
||||
return all(
|
||||
is_instance(k, key_type) and is_instance(v, val_type)
|
||||
for k, v in value.items()
|
||||
)
|
||||
|
||||
try:
|
||||
return isinstance(value, type_)
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
def from_dict(data_class, data):
|
||||
|
||||
init_values: MutableMapping[str, Any] = {}
|
||||
post_init_values: MutableMapping[str, Any] = {}
|
||||
|
||||
data_class_hints = get_type_hints(data_class)
|
||||
data_class_fields = cache(get_fields)(data_class)
|
||||
|
||||
extra_fields = set(data.keys()) - {f.name for f in data_class_fields}
|
||||
if extra_fields:
|
||||
formatted_keys = ", ".join(f'"{key}"' for key in extra_fields)
|
||||
raise ValueError(f"cannot match {formatted_keys} to any data class field")
|
||||
|
||||
for field in data_class_fields:
|
||||
field_type = data_class_hints[field.name]
|
||||
key = field.name
|
||||
|
||||
if key in data:
|
||||
try:
|
||||
value = _build_value(type_=field_type, data=data[key])
|
||||
except Exception as error:
|
||||
raise ValueError(error)
|
||||
|
||||
if not is_instance(value, field_type):
|
||||
raise ValueError((
|
||||
f'wrong value type for field "{field.name}" - should be "{field_type}" '
|
||||
f'instead of value "{value}" of type "{type(value)}"'
|
||||
))
|
||||
|
||||
init_values[field.name] = value
|
||||
|
||||
instance = data_class(**init_values)
|
||||
|
||||
for key, value in post_init_values.items():
|
||||
setattr(instance, key, value)
|
||||
|
||||
return instance
|
||||
|
||||
def normalize_chinese_punctuation(text):
|
||||
"""
|
||||
Convert Chinese (full-width) punctuation marks to English (half-width) equivalents.
|
||||
"""
|
||||
# Mapping of Chinese punctuation to English punctuation
|
||||
chinese_to_english_punct = {
|
||||
"，": ", ", # comma
|
||||
"。": ".", # period
|
||||
"：": ":", # colon
|
||||
"；": ";", # semicolon
|
||||
"？": "?", # question mark
|
||||
"！": "!", # exclamation mark
|
||||
"（": "(", # left parenthesis
|
||||
"）": ")", # right parenthesis
|
||||
"【": "[", # left square bracket
|
||||
"】": "]", # right square bracket
|
||||
"《": "<", # left angle quote
|
||||
"》": ">", # right angle quote
|
||||
"“": '"', # left double quotation
|
||||
"”": '"', # right double quotation
|
||||
"‘": "'", # left single quotation
|
||||
"’": "'", # right single quotation
|
||||
"、": ",", # enumeration comma
|
||||
"—": "-", # em dash
|
||||
"…": "...", # ellipsis
|
||||
"·": ".", # middle dot
|
||||
"「": '"', # left corner bracket
|
||||
"」": '"', # right corner bracket
|
||||
"『": '"', # left double corner bracket
|
||||
"』": '"', # right double corner bracket
|
||||
}
|
||||
|
||||
# Replace each Chinese punctuation with its English counterpart
|
||||
for zh_punct, en_punct in chinese_to_english_punct.items():
|
||||
text = text.replace(zh_punct, en_punct)
|
||||
|
||||
return text
|
||||
|
||||
def transcript_normalize(text: str):
|
||||
transcript = normalize_chinese_punctuation(text)
|
||||
|
||||
transcript = transcript.replace("(", " ")
|
||||
transcript = transcript.replace(")", " ")
|
||||
transcript = transcript.replace("°F", " degrees Fahrenheit")
|
||||
transcript = transcript.replace("°C", " degrees Celsius")
|
||||
|
||||
for tag, replacement in [
|
||||
("[laugh]", "<SE>[Laughter]</SE>"),
|
||||
("[humming start]", "<SE_s>[Humming]</SE_s>"),
|
||||
("[humming end]", "<SE_e>[Humming]</SE_e>"),
|
||||
("[music start]", "<SE_s>[Music]</SE_s>"),
|
||||
("[music end]", "<SE_e>[Music]</SE_e>"),
|
||||
("[music]", "<SE>[Music]</SE>"),
|
||||
("[sing start]", "<SE_s>[Singing]</SE_s>"),
|
||||
("[sing end]", "<SE_e>[Singing]</SE_e>"),
|
||||
("[applause]", "<SE>[Applause]</SE>"),
|
||||
("[cheering]", "<SE>[Cheering]</SE>"),
|
||||
("[cough]", "<SE>[Cough]</SE>"),
|
||||
]:
|
||||
transcript = transcript.replace(tag, replacement)
|
||||
lines = transcript.split("\n")
|
||||
transcript = "\n".join([" ".join(line.split()) for line in lines if line.strip()])
|
||||
transcript = transcript.strip()
|
||||
|
||||
if not any([transcript.endswith(c) for c in [".", "!", "?", ",", ";", '"', "'", "</SE_e>", "</SE>"]]):
|
||||
transcript += "."
|
||||
|
||||
return transcript
|
||||
|
||||
@dataclass
|
||||
class AudioContent:
|
||||
audio_url: str
|
||||
raw_audio: Optional[str] = None
|
||||
offset: Optional[float] = None
|
||||
duration: Optional[float] = None
|
||||
row_id: Optional[int] = None
|
||||
type: str = "audio"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextContent:
|
||||
text: str
|
||||
type: str = "text"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
role: str
|
||||
content: Union[str, AudioContent, TextContent, List[Union[str, AudioContent, TextContent]]]
|
||||
recipient: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatMLSample:
|
||||
messages: List[Message]
|
||||
start_index: Optional[int] = None
|
||||
misc: Optional[Dict] = None
|
||||
speaker: Optional[str] = None
|
||||
|
||||
def prepare_chatml_sample(sample: Union[ChatMLSample, Dict], tokenizer):
|
||||
|
||||
try:
|
||||
if not isinstance(sample, ChatMLSample):
|
||||
|
||||
# replacing pd.isna
|
||||
def is_nan(x):
|
||||
if isinstance(x, float):
|
||||
return math.isnan(x)
|
||||
if isinstance(x, np.generic):
|
||||
return np.isnan(x)
|
||||
if isinstance(x, torch.Tensor) and x.numel() == 1:
|
||||
return torch.isnan(x).item()
|
||||
return False
|
||||
|
||||
if "speaker" in sample and is_nan(sample["speaker"]):
|
||||
sample["speaker"] = None
|
||||
if "start_index" in sample and is_nan(sample["start_index"]):
|
||||
sample["start_index"] = None
|
||||
if "content" in sample and is_nan(sample["content"]):
|
||||
sample["content"] = ""
|
||||
|
||||
def convert_nan_to_none(obj):
|
||||
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
elif isinstance(obj, float) and math.isnan(obj):
|
||||
return None
|
||||
elif isinstance(obj, dict):
|
||||
return {k: convert_nan_to_none(v) for k, v in obj.items()}
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
return [convert_nan_to_none(item) for item in obj]
|
||||
return obj
|
||||
|
||||
clean_sample = convert_nan_to_none(sample)
|
||||
|
||||
val_keys = []
|
||||
for field in fields(ChatMLSample):
|
||||
if field.name in clean_sample:
|
||||
val_keys.append(field.name)
|
||||
clean_sample = {k: clean_sample[k] for k in val_keys}
|
||||
|
||||
sample = from_dict(
|
||||
data_class=ChatMLSample, data=clean_sample,
|
||||
)
|
||||
|
||||
input_tokens = []
|
||||
audio_contents = []
|
||||
speaker_id = None
|
||||
if sample.speaker is not None:
|
||||
speaker_id = sample.speaker
|
||||
elif sample.misc is not None:
|
||||
if "speaker" in sample.misc:
|
||||
speaker_id = sample.misc["speaker"]
|
||||
|
||||
total_m = len(sample.messages)
|
||||
for turn_id, message in enumerate(sample.messages):
|
||||
role = message.role
|
||||
recipient = message.recipient
|
||||
content = message.content
|
||||
content_l = []
|
||||
|
||||
if isinstance(content, str):
|
||||
content_l.append(TextContent(text=content))
|
||||
elif isinstance(content, TextContent):
|
||||
content_l.append(content)
|
||||
elif isinstance(content, AudioContent):
|
||||
content_l.append(content)
|
||||
elif isinstance(content, list):
|
||||
for ele in content:
|
||||
if isinstance(ele, str):
|
||||
content_l.append(TextContent(text=ele))
|
||||
else:
|
||||
content_l.append(ele)
|
||||
if turn_id == 0:
|
||||
prefix = f"<|begin_of_text|><|start_header_id|>{role}<|end_header_id|>\n\n"
|
||||
else:
|
||||
prefix = f"<|start_header_id|>{role}<|end_header_id|>\n\n"
|
||||
eot_postfix = "<|eot_id|>"
|
||||
eom_postfix = "<|eom_id|>"
|
||||
|
||||
prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
|
||||
input_tokens.extend(prefix_tokens)
|
||||
|
||||
if recipient:
|
||||
assert role == "assistant", "Recipient is only available for assistant role."
|
||||
recipient_tokens = tokenizer.encode(f"{recipient}<|recipient|>", add_special_tokens=False)
|
||||
input_tokens.extend(recipient_tokens)
|
||||
|
||||
for content in content_l:
|
||||
if content.type == "text":
|
||||
text_tokens = tokenizer.encode(content.text, add_special_tokens=False)
|
||||
input_tokens.extend(text_tokens)
|
||||
|
||||
elif content.type == "audio":
|
||||
audio_contents.append(content)
|
||||
if role == "user" or role == "system":
|
||||
text_tokens = tokenizer.encode(
|
||||
"<|audio_bos|><|AUDIO|><|audio_eos|>",
|
||||
add_special_tokens=False,
|
||||
)
|
||||
input_tokens.extend(text_tokens)
|
||||
elif role == "assistant":
|
||||
text_tokens = tokenizer.encode(
|
||||
"<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>",
|
||||
add_special_tokens=False,
|
||||
)
|
||||
input_tokens.extend(text_tokens)
|
||||
|
||||
next_id = turn_id + 1
|
||||
if role == "assistant" and next_id != total_m and sample.messages[next_id].role == "assistant":
|
||||
postfix_tokens = tokenizer.encode(eom_postfix, add_special_tokens=False)
|
||||
input_tokens.extend(postfix_tokens)
|
||||
else:
|
||||
postfix_tokens = tokenizer.encode(eot_postfix, add_special_tokens=False)
|
||||
input_tokens.extend(postfix_tokens)
|
||||
|
||||
return input_tokens, audio_contents, speaker_id
|
||||
|
||||
except Exception:
|
||||
return None, None, None
|
||||
|
||||
@dataclass
|
||||
class HiggsAudioBatchInput:
|
||||
input_ids: torch.LongTensor # shape (bsz, seq_len).
|
||||
attention_mask: torch.Tensor # shape (bsz, seq_len).
|
||||
audio_out_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length)
|
||||
audio_out_ids_start: Optional[torch.LongTensor] # shape (num_audio_out,)
|
||||
audio_out_ids_start_group_loc: Optional[torch.LongTensor] # shape (num_audio_out,), specify which a sample's group location in the batch
|
||||
audio_in_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_in_total_length)
|
||||
audio_in_ids_start: Optional[torch.LongTensor] # shape (num_audio_in,)
|
||||
label_ids: Optional[torch.LongTensor] # shape (bsz, seq_len)
|
||||
label_audio_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length)
|
||||
reward: Optional[float] = None
|
||||
|
||||
@dataclass
|
||||
class ChatMLDatasetSample:
|
||||
|
||||
input_ids: torch.LongTensor # (seq_len,) Input text tokens
|
||||
label_ids: torch.LongTensor # (seq_len,) Label IDs
|
||||
audio_ids_concat: torch.LongTensor # (num_codebooks, audio_seq_len) Concatenated audio tokens
|
||||
audio_ids_start: torch.LongTensor # (num_audios,) Start index of each audio token in `audio_ids_concat`
|
||||
audio_waveforms_concat: torch.Tensor # (total_wv_length,) Concatenated audio waveforms
|
||||
audio_waveforms_start: torch.LongTensor # (num_audios,) Start index of each waveform in `audio_waveforms_concat`
|
||||
audio_sample_rate: torch.Tensor # (num_audios,) Sampling rate per audio waveform
|
||||
audio_speaker_indices: torch.LongTensor # (num_audios,) Speaker indices per audio; -1 = unknown
|
||||
audio_label_ids_concat: Optional[torch.LongTensor] = None # (num_codebooks, audio_seq_len) Optional audio token labels
|
||||
reward: Optional[float] = None # Optional scalar reward
|
||||
|
||||
|
||||
def num_audios(self):
|
||||
return max(len(self.audio_waveforms_start), len(self.audio_ids_start))
|
||||
|
||||
def get_audio_codes(self, idx):
|
||||
code_start = self.audio_ids_start[idx]
|
||||
if idx < len(self.audio_ids_start) - 1:
|
||||
code_end = self.audio_ids_start[idx + 1]
|
||||
else:
|
||||
code_end = self.audio_ids_concat.shape[-1]
|
||||
|
||||
return self.audio_ids_concat[:, code_start:code_end]
|
||||
|
||||
def get_audio_codes_labels(self, idx):
|
||||
if self.audio_label_ids_concat is None:
|
||||
return None
|
||||
code_start = self.audio_ids_start[idx]
|
||||
if idx < len(self.audio_ids_start) - 1:
|
||||
code_end = self.audio_ids_start[idx + 1]
|
||||
else:
|
||||
code_end = self.audio_ids_concat.shape[-1]
|
||||
|
||||
return self.audio_label_ids_concat[:, code_start:code_end]
|
||||
|
||||
def get_wv(self, idx):
|
||||
wv_start = self.audio_waveforms_start[idx]
|
||||
sr = self.audio_sample_rate[idx]
|
||||
if idx < len(self.audio_waveforms_start) - 1:
|
||||
wv_end = self.audio_waveforms_start[idx + 1]
|
||||
else:
|
||||
wv_end = self.audio_waveforms_concat.shape[-1]
|
||||
return self.audio_waveforms_concat[wv_start:wv_end], sr
|
||||
|
||||
class HiggsAudioSampleCollator:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
audio_in_token_id,
|
||||
audio_out_token_id,
|
||||
pad_token_id,
|
||||
audio_stream_bos_id,
|
||||
audio_stream_eos_id,
|
||||
round_to=8,
|
||||
pad_left=False,
|
||||
return_audio_in_tokens=True,
|
||||
audio_num_codebooks=None,
|
||||
use_delay_pattern=False,
|
||||
disable_audio_codes_transform=False,
|
||||
add_new_bos_eos_for_long_chunk=True,
|
||||
mask_audio_out_token_label=True,
|
||||
):
|
||||
self.round_to = round_to
|
||||
self.pad_left = pad_left
|
||||
self.audio_in_token_id = audio_in_token_id
|
||||
self.audio_out_token_id = audio_out_token_id
|
||||
self.audio_stream_bos_id = audio_stream_bos_id
|
||||
self.audio_stream_eos_id = audio_stream_eos_id
|
||||
self.pad_token_id = pad_token_id
|
||||
self.return_audio_in_tokens = return_audio_in_tokens
|
||||
self.audio_num_codebooks = audio_num_codebooks
|
||||
self.use_delay_pattern = use_delay_pattern
|
||||
|
||||
self.disable_audio_codes_transform = disable_audio_codes_transform
|
||||
self.add_new_bos_eos_for_long_chunk = add_new_bos_eos_for_long_chunk
|
||||
self.mask_audio_out_token_label = mask_audio_out_token_label
|
||||
|
||||
def _process_and_duplicate_audio_tokens(
|
||||
self, input_ids: torch.Tensor, audio_idx: int, wv: torch.Tensor, labels: Optional[torch.Tensor] = None
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, int]:
|
||||
|
||||
total_samples = len(wv)
|
||||
num_chunks = math.ceil(total_samples / self.chunk_size_samples)
|
||||
|
||||
if num_chunks <= 1:
|
||||
return input_ids, labels, 1
|
||||
|
||||
audio_token_seq = input_ids[audio_idx - 1 : audio_idx + 2]
|
||||
duplicated_sequence = audio_token_seq.repeat(num_chunks)
|
||||
|
||||
new_input_ids = torch.cat([input_ids[: audio_idx - 1], duplicated_sequence, input_ids[audio_idx + 2 :]])
|
||||
|
||||
new_labels = None
|
||||
if labels is not None:
|
||||
label_seq = labels[audio_idx - 1 : audio_idx + 2]
|
||||
duplicated_labels = label_seq.repeat(num_chunks)
|
||||
new_labels = torch.cat([labels[: audio_idx - 1], duplicated_labels, labels[audio_idx + 2 :]])
|
||||
|
||||
return new_input_ids, new_labels, num_chunks
|
||||
|
||||
def __call__(self, batch: List[ChatMLDatasetSample]):
|
||||
|
||||
label_ids = None
|
||||
label_audio_ids = None
|
||||
if all([ele.label_ids is None for ele in batch]):
|
||||
return_labels = False
|
||||
else:
|
||||
return_labels = True
|
||||
|
||||
processed_batch = batch
|
||||
|
||||
# Get the max sequence length based on processed batch
|
||||
max_seq_length = _ceil_to_nearest(max([len(sample.input_ids) for sample in processed_batch]), self.round_to)
|
||||
|
||||
# Get the ids for audio-in and audio-out for each batch
|
||||
audio_in_ids_l = []
|
||||
audio_out_ids_l = []
|
||||
audio_out_ids_group_loc_l = []
|
||||
audio_in_label_ids_l = None
|
||||
audio_out_label_ids_l = None
|
||||
reward_l = []
|
||||
|
||||
if return_labels:
|
||||
audio_out_no_train_flag = [] # Whether the audio-out data should be trained on or not.
|
||||
|
||||
# Process the audio inputs and outputs
|
||||
for i in range(len(processed_batch)):
|
||||
audio_in_mask = processed_batch[i].input_ids == self.audio_in_token_id
|
||||
audio_out_mask = processed_batch[i].input_ids == self.audio_out_token_id
|
||||
audio_ids = torch.ones_like(processed_batch[i].input_ids)
|
||||
audio_ids[audio_in_mask ^ audio_out_mask] = torch.cumsum(audio_ids[audio_in_mask ^ audio_out_mask], 0) - 1
|
||||
audio_in_ids = audio_ids[audio_in_mask]
|
||||
audio_out_ids = audio_ids[audio_out_mask]
|
||||
|
||||
if return_labels:
|
||||
audio_out_no_train_flag.append(processed_batch[i].label_ids[audio_out_mask] < 0)
|
||||
if self.mask_audio_out_token_label:
|
||||
processed_batch[i].label_ids[audio_out_mask] = -100
|
||||
|
||||
if self.return_audio_in_tokens:
|
||||
audio_in_ids_l.extend(
|
||||
[processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_in_ids]
|
||||
)
|
||||
if processed_batch[i].audio_label_ids_concat is not None:
|
||||
if audio_in_label_ids_l is None:
|
||||
audio_in_label_ids_l = []
|
||||
audio_in_label_ids_l.extend(
|
||||
[
|
||||
processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]
|
||||
for idx in audio_in_ids
|
||||
]
|
||||
)
|
||||
|
||||
audio_out_ids_l.extend(
|
||||
[processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_out_ids]
|
||||
)
|
||||
audio_out_ids_group_loc_l.append(i)
|
||||
if processed_batch[i].reward is not None:
|
||||
reward_l.append(processed_batch[i].reward)
|
||||
|
||||
if processed_batch[i].audio_label_ids_concat is not None:
|
||||
if audio_out_label_ids_l is None:
|
||||
audio_out_label_ids_l = []
|
||||
audio_out_label_ids_l.extend(
|
||||
[
|
||||
processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]
|
||||
for idx in audio_out_ids
|
||||
]
|
||||
)
|
||||
|
||||
if return_labels:
|
||||
audio_out_no_train_flag = torch.cat(audio_out_no_train_flag, dim=0)
|
||||
|
||||
if len(audio_in_ids_l) > 0:
|
||||
|
||||
# I tried to remove the for-loop in original implementation
|
||||
# but to do batching with padding caused problem so I turned it into a list compre.
|
||||
lengths = [seg.shape[1] for seg in audio_in_ids_l]
|
||||
aug_lengths = [length + 2 for length in lengths]
|
||||
audio_in_ids_start = torch.cumsum(
|
||||
torch.tensor([0] + aug_lengths[:-1], dtype=torch.long), dim=0
|
||||
)
|
||||
|
||||
if self.disable_audio_codes_transform:
|
||||
audio_in_ids = torch.cat(audio_in_ids_l, dim=1).long()
|
||||
else:
|
||||
with_tokens = [
|
||||
torch.cat([
|
||||
torch.full((seg.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),
|
||||
seg,
|
||||
torch.full((seg.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
|
||||
], dim=1)
|
||||
for seg in audio_in_ids_l
|
||||
]
|
||||
|
||||
if self.use_delay_pattern:
|
||||
with_tokens = [
|
||||
build_delay_pattern_mask(
|
||||
tok.unsqueeze(0),
|
||||
bos_token_id=self.audio_stream_bos_id,
|
||||
pad_token_id=self.audio_stream_eos_id
|
||||
)[0]
|
||||
for tok in with_tokens
|
||||
]
|
||||
audio_in_ids = torch.cat(with_tokens, dim=1).long()
|
||||
else:
|
||||
audio_in_ids = torch.zeros((0, 0), dtype=torch.long)
|
||||
audio_in_ids_start = torch.zeros(0, dtype=torch.long)
|
||||
|
||||
audio_out_ids_start_group_loc = None
|
||||
if len(audio_out_ids_l) > 0:
|
||||
new_audio_out_ids_l = []
|
||||
label_audio_ids_l = []
|
||||
for idx, ele in enumerate(audio_out_ids_l):
|
||||
if self.disable_audio_codes_transform:
|
||||
|
||||
audio_codes = ele
|
||||
if return_labels:
|
||||
label_audio_ids = audio_out_label_ids_l[idx]
|
||||
else:
|
||||
audio_codes = torch.cat(
|
||||
[
|
||||
torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),
|
||||
ele,
|
||||
torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
if return_labels:
|
||||
label_audio_ids = torch.cat(
|
||||
[
|
||||
torch.full((ele.shape[0], 1), -100, dtype=torch.long),
|
||||
ele,
|
||||
torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
if self.use_delay_pattern:
|
||||
audio_codes = build_delay_pattern_mask(
|
||||
audio_codes.unsqueeze(0),
|
||||
bos_token_id=self.audio_stream_bos_id,
|
||||
pad_token_id=self.audio_stream_eos_id,
|
||||
)[0].squeeze(0)
|
||||
if return_labels:
|
||||
label_audio_ids = build_delay_pattern_mask(
|
||||
label_audio_ids.unsqueeze(0),
|
||||
bos_token_id=-100,
|
||||
pad_token_id=-100,
|
||||
)[0].squeeze(0)
|
||||
new_audio_out_ids_l.append(audio_codes)
|
||||
|
||||
if return_labels:
|
||||
if audio_out_no_train_flag[idx]:
|
||||
label_audio_ids[:] = -100
|
||||
label_audio_ids_l.append(label_audio_ids)
|
||||
|
||||
audio_out_ids = torch.cat(new_audio_out_ids_l, dim=1).long()
|
||||
if return_labels:
|
||||
label_audio_ids = torch.cat(label_audio_ids_l, dim=1).long()
|
||||
audio_out_ids_start = torch.cumsum(
|
||||
torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_out_ids_l[:-1]]), dim=0
|
||||
)
|
||||
audio_out_ids_start_group_loc = torch.tensor(audio_out_ids_group_loc_l, dtype=torch.long)
|
||||
else:
|
||||
audio_out_ids = torch.zeros((0, 0), dtype=torch.long)
|
||||
audio_out_ids_start = torch.zeros(0, dtype=torch.long)
|
||||
if return_labels:
|
||||
label_audio_ids = torch.zeros((0, 0), dtype=torch.long)
|
||||
|
||||
reward = torch.tensor(reward_l, dtype=torch.float32)
|
||||
|
||||
# cleaner and faster implementation
|
||||
def pad_sequence(seq, length, pad_value):
|
||||
if self.pad_left:
|
||||
padding = (length - len(seq), 0)
|
||||
else:
|
||||
padding = (0, length - len(seq))
|
||||
return F.pad(seq, padding, value=pad_value)
|
||||
|
||||
input_ids = torch.stack([
|
||||
pad_sequence(ele.input_ids, max_seq_length, self.pad_token_id)
|
||||
for ele in processed_batch
|
||||
])
|
||||
|
||||
if return_labels:
|
||||
label_ids = torch.stack([
|
||||
pad_sequence(ele.label_ids, max_seq_length, -100)
|
||||
for ele in processed_batch
|
||||
])
|
||||
|
||||
attention_mask = torch.stack([
|
||||
pad_sequence(torch.ones_like(ele.input_ids), max_seq_length, 0)
|
||||
for ele in processed_batch
|
||||
])
|
||||
|
||||
if not self.return_audio_in_tokens:
|
||||
audio_in_ids = None
|
||||
audio_in_ids_start = None
|
||||
|
||||
if self.audio_num_codebooks is not None:
|
||||
if audio_in_ids is not None:
|
||||
audio_in_ids = audio_in_ids[: self.audio_num_codebooks]
|
||||
if audio_out_ids is not None:
|
||||
audio_out_ids = audio_out_ids[: self.audio_num_codebooks]
|
||||
if label_audio_ids is not None:
|
||||
label_audio_ids = label_audio_ids[: self.audio_num_codebooks]
|
||||
|
||||
return HiggsAudioBatchInput(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
audio_out_ids=audio_out_ids,
|
||||
audio_out_ids_start=audio_out_ids_start,
|
||||
audio_out_ids_start_group_loc=audio_out_ids_start_group_loc,
|
||||
audio_in_ids=audio_in_ids,
|
||||
audio_in_ids_start=audio_in_ids_start,
|
||||
label_ids=label_ids,
|
||||
label_audio_ids=label_audio_ids,
|
||||
reward=reward,
|
||||
)
|
||||
871
comfy/ldm/higgsv2/tokenizer.py
Normal file
871
comfy/ldm/higgsv2/tokenizer.py
Normal file
@ -0,0 +1,871 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from typing import Optional
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils.parametrizations import weight_norm
|
||||
import torchaudio
|
||||
import numpy as np
|
||||
from torch import vmap
|
||||
from transformers import AutoModel
|
||||
|
||||
def WNConv1d(*args, device = None, dtype = None, operations = None, **kwargs):
|
||||
return weight_norm(operations.Conv1d(*args, **kwargs, device = device, dtype = dtype))
|
||||
|
||||
|
||||
def WNConvTranspose1d(*args, device = None, dtype = None, operations = None, **kwargs):
|
||||
return weight_norm(operations.ConvTranspose1d(*args, **kwargs, device = device, dtype = dtype))
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def snake(x, alpha):
|
||||
shape = x.shape
|
||||
x = x.reshape(shape[0], shape[1], -1)
|
||||
x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
|
||||
x = x.reshape(shape)
|
||||
return x
|
||||
|
||||
|
||||
class Snake1d(nn.Module):
|
||||
def __init__(self, channels, device = None, dtype = None):
|
||||
super().__init__()
|
||||
self.alpha = nn.Parameter(torch.ones(1, channels, 1, device = device, dtype = dtype))
|
||||
|
||||
def forward(self, x):
|
||||
return snake(x, self.alpha)
|
||||
|
||||
class DACResidualUnit(nn.Module):
|
||||
def __init__(self, dim: int = 16, dilation: int = 1, device = None, dtype = None, operations = None):
|
||||
super().__init__()
|
||||
pad = ((7 - 1) * dilation) // 2
|
||||
self.block = nn.Sequential(
|
||||
Snake1d(dim, device = device, dtype = dtype),
|
||||
WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad, device = device, dtype = dtype, operations = operations),
|
||||
Snake1d(dim, device = device, dtype = dtype),
|
||||
WNConv1d(dim, dim, kernel_size=1, device = device, dtype = dtype, operations = operations),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.block(x)
|
||||
pad = (x.shape[-1] - y.shape[-1]) // 2
|
||||
if pad > 0:
|
||||
x = x[..., pad:-pad]
|
||||
return x + y
|
||||
|
||||
|
||||
class DACEncoderBlock(nn.Module):
|
||||
def __init__(self, dim: int = 16, stride: int = 1, device = None, dtype = None, operations = None):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
DACResidualUnit(dim // 2, dilation=1, device = device, dtype = dtype, operations = operations),
|
||||
DACResidualUnit(dim // 2, dilation=3, device = device, dtype = dtype, operations = operations),
|
||||
DACResidualUnit(dim // 2, dilation=9, device = device, dtype = dtype, operations = operations),
|
||||
Snake1d(dim // 2),
|
||||
WNConv1d(
|
||||
dim // 2,
|
||||
dim,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
padding=math.ceil(stride / 2),
|
||||
device = device, dtype = dtype, operations = operations
|
||||
),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class DACEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int = 64,
|
||||
strides: list = [2, 4, 8, 8],
|
||||
d_latent: int = 256,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
# Create first convolution
|
||||
self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations)]
|
||||
|
||||
# Create EncoderBlocks that double channels as they downsample by `stride`
|
||||
for stride in strides:
|
||||
d_model *= 2
|
||||
self.block += [DACEncoderBlock(d_model, stride=stride, device = device, dtype = dtype, operations = operations)]
|
||||
|
||||
# Create last convolution
|
||||
self.block += [
|
||||
Snake1d(d_model),
|
||||
WNConv1d(d_model, d_latent, kernel_size=3, padding=1, device = device, dtype = dtype, operations = operations),
|
||||
]
|
||||
|
||||
# Wrap black into nn.Sequential
|
||||
self.block = nn.Sequential(*self.block)
|
||||
self.enc_dim = d_model
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class DACDecoderBlock(nn.Module):
|
||||
def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, device = None, dtype = None, operations = None):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
Snake1d(input_dim, device = device, dtype = dtype),
|
||||
WNConvTranspose1d(
|
||||
input_dim,
|
||||
output_dim,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
padding=math.ceil(stride / 2),
|
||||
output_padding=stride % 2, # out_pad,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
),
|
||||
DACResidualUnit(output_dim, dilation=1, device = device, dtype = dtype, operations = operations),
|
||||
DACResidualUnit(output_dim, dilation=3, device = device, dtype = dtype, operations = operations),
|
||||
DACResidualUnit(output_dim, dilation=9, device = device, dtype = dtype, operations = operations),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class DACDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channel,
|
||||
channels,
|
||||
rates,
|
||||
d_out: int = 1,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Add first conv layer
|
||||
layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations )]
|
||||
|
||||
# Add upsampling + MRF blocks
|
||||
for i, stride in enumerate(rates):
|
||||
input_dim = channels // 2**i
|
||||
output_dim = channels // 2 ** (i + 1)
|
||||
layers += [DACDecoderBlock(input_dim, output_dim, stride, device = device, dtype = dtype, operations = operations)]
|
||||
|
||||
# Add final conv layer
|
||||
layers += [
|
||||
Snake1d(output_dim, device = device, dtype = dtype),
|
||||
WNConv1d(output_dim, d_out, kernel_size=7, padding=3, device = device, dtype = dtype, operations = operations),
|
||||
]
|
||||
|
||||
self.model = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
class Conv1d1x1:
|
||||
def __new__(cls, in_channels, out_channels, bias=True, device=None, dtype=None, operations=None):
|
||||
operations = operations or nn
|
||||
return operations.Conv1d(
|
||||
in_channels, out_channels, kernel_size=1,
|
||||
bias=bias, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
class Conv1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
padding: int = -1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if padding < 0:
|
||||
padding = (kernel_size - 1) // 2 * dilation
|
||||
|
||||
self.dilation = dilation
|
||||
self.conv = operations.Conv1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
class ConvTranspose1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int,
|
||||
padding=-1,
|
||||
output_padding=-1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
if padding < 0:
|
||||
padding = (stride + 1) // 2
|
||||
if output_padding < 0:
|
||||
output_padding = 1 if stride % 2 else 0
|
||||
self.deconv = operations.ConvTranspose1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
output_padding=output_padding,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.deconv(x)
|
||||
return x
|
||||
|
||||
class ResidualUnit(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
nonlinear_activation="ELU",
|
||||
nonlinear_activation_params={},
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
self.activation = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
|
||||
self.conv1 = Conv1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=1,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
self.conv2 = Conv1d1x1(out_channels, out_channels, bias, device = device, dtype = dtype, operations = operations)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.conv1(self.activation(x))
|
||||
y = self.conv2(self.activation(y))
|
||||
return x + y
|
||||
|
||||
|
||||
class EncoderBlock(nn.Module):
|
||||
def __init__(
|
||||
self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True, device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
self.res_units = torch.nn.ModuleList()
|
||||
for dilation in dilations:
|
||||
self.res_units += [ResidualUnit(in_channels, in_channels, kernel_size=unit_kernel_size, dilation=dilation, device = device, dtype = dtype, operations = operations)]
|
||||
self.num_res = len(self.res_units)
|
||||
|
||||
kernel_size=3 if stride == 1 else (2 * stride) # special case: stride=1, do not use kernel=2
|
||||
self.conv = Conv1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size = kernel_size,
|
||||
stride=stride,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
for idx in range(self.num_res):
|
||||
x = self.res_units[idx](x)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
encode_channels: int,
|
||||
channel_ratios=(1, 1),
|
||||
strides=(1, 1),
|
||||
kernel_size=3,
|
||||
bias=True,
|
||||
block_dilations=(1, 1),
|
||||
unit_kernel_size=3,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
assert len(channel_ratios) == len(strides)
|
||||
self.conv = Conv1d(
|
||||
in_channels=input_channels, out_channels=encode_channels, kernel_size=kernel_size, stride=1, bias=False,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
self.conv_blocks = torch.nn.ModuleList()
|
||||
in_channels = encode_channels
|
||||
for idx, stride in enumerate(strides):
|
||||
out_channels = int(encode_channels * channel_ratios[idx]) # could be float
|
||||
self.conv_blocks += [
|
||||
EncoderBlock(
|
||||
in_channels,
|
||||
out_channels,
|
||||
stride,
|
||||
dilations=block_dilations,
|
||||
unit_kernel_size=unit_kernel_size,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
]
|
||||
in_channels = out_channels
|
||||
self.num_blocks = len(self.conv_blocks)
|
||||
self.out_channels = out_channels
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
for i in range(self.num_blocks):
|
||||
x = self.conv_blocks[i](x)
|
||||
return x
|
||||
|
||||
|
||||
class DecoderBlock(nn.Module):
|
||||
"""Decoder block (no up-sampling)"""
|
||||
|
||||
def __init__(
|
||||
self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True, device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if stride == 1:
|
||||
self.conv = Conv1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=3, # fix kernel=3 when stride=1 for unchanged shape
|
||||
stride=stride,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
else:
|
||||
self.conv = ConvTranspose1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=(2 * stride),
|
||||
stride=stride,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
self.res_units = nn.ModuleList([
|
||||
ResidualUnit(out_channels, out_channels, kernel_size=unit_kernel_size, dilation=d, device = device, dtype = dtype, operations = operations)
|
||||
for d in dilations
|
||||
])
|
||||
|
||||
self.num_res = len(self.res_units)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
for idx in range(self.num_res):
|
||||
x = self.res_units[idx](x)
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
code_dim: int,
|
||||
output_channels: int,
|
||||
decode_channels: int,
|
||||
channel_ratios=(1, 1),
|
||||
strides=(1, 1),
|
||||
kernel_size=3,
|
||||
bias=True,
|
||||
block_dilations=(1, 1),
|
||||
unit_kernel_size=3,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
assert len(channel_ratios) == len(strides)
|
||||
self.conv1 = Conv1d(
|
||||
in_channels=code_dim,
|
||||
out_channels=int(decode_channels * channel_ratios[0]),
|
||||
kernel_size=kernel_size,
|
||||
stride=1,
|
||||
bias=False,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
self.conv_blocks = torch.nn.ModuleList()
|
||||
for idx, stride in enumerate(strides):
|
||||
in_channels = int(decode_channels * channel_ratios[idx])
|
||||
if idx < (len(channel_ratios) - 1):
|
||||
out_channels = int(decode_channels * channel_ratios[idx + 1])
|
||||
else:
|
||||
out_channels = decode_channels
|
||||
self.conv_blocks += [
|
||||
DecoderBlock(
|
||||
in_channels,
|
||||
out_channels,
|
||||
stride,
|
||||
dilations=block_dilations,
|
||||
unit_kernel_size=unit_kernel_size,
|
||||
bias=bias,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
]
|
||||
self.num_blocks = len(self.conv_blocks)
|
||||
|
||||
self.conv2 = Conv1d(out_channels, output_channels, kernel_size = 3, bias=False, device = device, dtype = dtype, operations = operations)
|
||||
|
||||
def forward(self, z):
|
||||
x = self.conv1(z)
|
||||
for i in range(self.num_blocks):
|
||||
x = self.conv_blocks[i](x)
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
class HiggsAudioFeatureExtractor(nn.Module):
|
||||
def __init__(self, sampling_rate=16000):
|
||||
super().__init__()
|
||||
self.sampling_rate = sampling_rate
|
||||
|
||||
def forward(self, audio_signal):
|
||||
audio_signal = audio_signal.unsqueeze(0)
|
||||
if len(audio_signal.shape) < 3:
|
||||
audio_signal = audio_signal.unsqueeze(0)
|
||||
return {"input_values": audio_signal}
|
||||
|
||||
def uniform_init(*shape: int, device = None, dtype = None):
|
||||
t = torch.empty(shape, device = device, dtype = dtype)
|
||||
nn.init.kaiming_uniform_(t)
|
||||
return t
|
||||
|
||||
class EuclideanCodebook(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
codebook_size: int,
|
||||
kmeans_init: int = False,
|
||||
kmeans_iters: int = 10,
|
||||
decay: float = 0.99,
|
||||
epsilon: float = 1e-5,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
device = None, dtype = None
|
||||
):
|
||||
super().__init__()
|
||||
self.decay = decay
|
||||
init_fn = uniform_init
|
||||
embed = init_fn(codebook_size, dim, device = device, dtype = dtype)
|
||||
|
||||
self.codebook_size = codebook_size
|
||||
|
||||
self.kmeans_iters = kmeans_iters
|
||||
self.epsilon = epsilon
|
||||
self.threshold_ema_dead_code = threshold_ema_dead_code
|
||||
|
||||
# Flag variable to indicate whether the codebook is initialized
|
||||
self.register_buffer("inited", torch.Tensor([not kmeans_init]))
|
||||
# Runing EMA cluster size/count: N_i^t in eq. (6) in vqvae paper
|
||||
self.register_buffer("cluster_size", torch.zeros(codebook_size))
|
||||
# Codebook
|
||||
self.register_buffer("embed", embed)
|
||||
# EMA codebook: eq. (7) in vqvae paper
|
||||
self.register_buffer("embed_avg", embed.clone())
|
||||
|
||||
def preprocess(self, x):
|
||||
x = x.view(-1, x.shape[-1])
|
||||
return x
|
||||
|
||||
def quantize(self, x):
|
||||
embed = self.embed.t()
|
||||
if x.dtype != embed.dtype:
|
||||
x = x.to(embed.dtype)
|
||||
|
||||
dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))
|
||||
embed_ind = dist.max(dim=-1).indices
|
||||
return embed_ind
|
||||
|
||||
def postprocess_emb(self, embed_ind, shape):
|
||||
return embed_ind.view(*shape[:-1])
|
||||
|
||||
def dequantize(self, embed_ind):
|
||||
quantize = F.embedding(embed_ind, self.embed)
|
||||
return quantize
|
||||
|
||||
def encode(self, x):
|
||||
shape = x.shape
|
||||
# pre-process
|
||||
x = self.preprocess(x) # [B, T, D] -> [B*T, D]
|
||||
# quantize
|
||||
embed_ind = self.quantize(x)
|
||||
# post-process
|
||||
embed_ind = self.postprocess_emb(embed_ind, shape)
|
||||
return embed_ind
|
||||
|
||||
def decode(self, embed_ind):
|
||||
quantize = self.dequantize(embed_ind)
|
||||
return quantize
|
||||
|
||||
def forward(self, x):
|
||||
orig_shape = x.shape # [B, T, D]
|
||||
flat = x.view(-1, x.shape[-1]) # [B*T, D]
|
||||
|
||||
embed_ind = self.quantize(flat)
|
||||
embed_ind = self.postprocess_emb(embed_ind, orig_shape)
|
||||
# now embed_ind has shape [B, T]
|
||||
|
||||
quantize = self.dequantize(embed_ind)
|
||||
# quantize: [B, T, D]
|
||||
|
||||
return quantize, embed_ind
|
||||
|
||||
class VectorQuantization(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
codebook_size: int,
|
||||
codebook_dim: Optional[int] = None,
|
||||
decay: float = 0.99,
|
||||
epsilon: float = 1e-5,
|
||||
kmeans_init: bool = True,
|
||||
kmeans_iters: int = 50,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
commitment_weight: float = 1.0,
|
||||
device = None, dtype = None, operations = None
|
||||
):
|
||||
super().__init__()
|
||||
_codebook_dim: int = codebook_dim if codebook_dim is not None else dim
|
||||
|
||||
requires_projection = _codebook_dim != dim
|
||||
self.project_in = operations.Linear(dim, _codebook_dim, device = device, dtype = dtype) if requires_projection else nn.Identity()
|
||||
self.project_out = operations.Linear(_codebook_dim, dim, device = device, dtype = dtype) if requires_projection else nn.Identity()
|
||||
|
||||
self.epsilon = epsilon
|
||||
self.commitment_weight = commitment_weight
|
||||
|
||||
self._codebook = EuclideanCodebook(
|
||||
dim=_codebook_dim,
|
||||
codebook_size=codebook_size,
|
||||
kmeans_init=kmeans_init,
|
||||
kmeans_iters=kmeans_iters,
|
||||
decay=decay,
|
||||
epsilon=epsilon,
|
||||
threshold_ema_dead_code=threshold_ema_dead_code,
|
||||
device = device, dtype = dtype
|
||||
)
|
||||
self.codebook_size = codebook_size
|
||||
|
||||
@property
|
||||
def codebook(self):
|
||||
return self._codebook.embed
|
||||
|
||||
def encode(self, x):
|
||||
x = x.permute(0, 2, 1)
|
||||
x = self.project_in(x)
|
||||
embed_in = self._codebook.encode(x)
|
||||
return embed_in
|
||||
|
||||
def decode(self, embed_ind):
|
||||
quantize = self._codebook.decode(embed_ind)
|
||||
quantize = self.project_out(quantize)
|
||||
quantize = quantize.permute(0, 2, 1)
|
||||
return quantize
|
||||
|
||||
def forward(self, x):
|
||||
device = x.device
|
||||
x = x.transpose(1, 2).contiguous() # [b d n] -> [b n d]
|
||||
x = self.project_in(x)
|
||||
|
||||
quantize, embed_ind = self._codebook(x)
|
||||
|
||||
loss = torch.tensor([0.0], device=device, requires_grad=self.training)
|
||||
|
||||
quantize = self.project_out(quantize)
|
||||
quantize = quantize.transpose(1, 2).contiguous() # [b n d] -> [b d n]
|
||||
return quantize, embed_ind, loss
|
||||
|
||||
|
||||
class ResidualVectorQuantization(nn.Module):
|
||||
def __init__(self, *, num_quantizers, device = None, dtype = None, operations = None, **kwargs):
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleList([VectorQuantization(device = device, dtype = dtype, operations = operations, **kwargs) for _ in range(num_quantizers)])
|
||||
|
||||
def forward(self, x, n_q: Optional[int] = None):
|
||||
quantized_out = 0.0
|
||||
residual = x
|
||||
|
||||
all_losses = []
|
||||
all_indices = []
|
||||
|
||||
n_q = n_q or len(self.layers)
|
||||
|
||||
for layer in self.layers[:n_q]:
|
||||
quantized, indices, loss = layer(residual)
|
||||
residual = residual - quantized
|
||||
quantized_out = quantized_out + quantized
|
||||
|
||||
all_indices.append(indices)
|
||||
all_losses.append(loss)
|
||||
|
||||
out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
|
||||
return quantized_out, out_indices, out_losses
|
||||
|
||||
def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
|
||||
""" Vectorized Implementation of dequantization | 2x faster than original impl """
|
||||
|
||||
biases = torch.stack([layer.project_out.bias for layer in self.layers])
|
||||
|
||||
codebook_device = self.layers[0]._codebook.embed.device
|
||||
q_indices = q_indices.to(codebook_device)
|
||||
|
||||
def decode_one(codebook_weight, proj_weight, embed_id, proj_biases):
|
||||
quantized = F.embedding(embed_id, codebook_weight).transpose(1, 2) # (B, D, T)
|
||||
quantized = F.linear(quantized.transpose(1, 2), proj_weight, proj_biases).transpose(1, 2)
|
||||
return quantized
|
||||
|
||||
codebook_weights = torch.stack([q._codebook.embed for q in self.layers]) # (n_codebooks, vocab_size, D)
|
||||
proj_weights = torch.stack([q.project_out.weight for q in self.layers])
|
||||
|
||||
quantized = vmap(decode_one)(codebook_weights, proj_weights, q_indices, biases)
|
||||
|
||||
return quantized.sum(0)
|
||||
|
||||
class ResidualVectorQuantizer(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimension: int = 256,
|
||||
codebook_dim: int = None,
|
||||
n_q: int = 8,
|
||||
bins: int = 1024,
|
||||
decay: float = 0.99,
|
||||
kmeans_init: bool = True,
|
||||
kmeans_iters: int = 50,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
device = None,
|
||||
dtype = None,
|
||||
operations = None
|
||||
):
|
||||
super().__init__()
|
||||
self.n_q = n_q
|
||||
self.dimension = dimension
|
||||
self.codebook_dim = codebook_dim
|
||||
self.bins = bins
|
||||
self.decay = decay
|
||||
self.kmeans_init = kmeans_init
|
||||
self.kmeans_iters = kmeans_iters
|
||||
self.threshold_ema_dead_code = threshold_ema_dead_code
|
||||
self.vq = ResidualVectorQuantization(
|
||||
dim=self.dimension,
|
||||
codebook_dim=self.codebook_dim,
|
||||
codebook_size=self.bins,
|
||||
num_quantizers=self.n_q,
|
||||
decay=self.decay,
|
||||
kmeans_init=self.kmeans_init,
|
||||
kmeans_iters=self.kmeans_iters,
|
||||
threshold_ema_dead_code=self.threshold_ema_dead_code,
|
||||
device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, sample_rate: int, bandwidth: Optional[float] = None): # -> QuantizedResult:
|
||||
|
||||
bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)
|
||||
n_q = self.get_num_quantizers_for_bandwidth(sample_rate, bandwidth)
|
||||
quantized, codes, commit_loss = self.vq(x, n_q=n_q)
|
||||
bw = torch.tensor(n_q * bw_per_q).to(x)
|
||||
return quantized, codes, bw, torch.mean(commit_loss)
|
||||
|
||||
def get_num_quantizers_for_bandwidth(self, sample_rate: int, bandwidth: Optional[float] = None) -> int:
|
||||
"""Return n_q based on specified target bandwidth."""
|
||||
bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)
|
||||
n_q = self.n_q
|
||||
if bandwidth and bandwidth > 0.0:
|
||||
n_q = int(max(1, math.floor(bandwidth / bw_per_q)))
|
||||
return n_q
|
||||
|
||||
def get_bandwidth_per_quantizer(self, sample_rate: int):
|
||||
"""Return bandwidth per quantizer for a given input sample rate."""
|
||||
return math.log2(self.bins) * sample_rate / 1000
|
||||
|
||||
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
"""Decode the given codes to the quantized representation."""
|
||||
quantized = self.vq.decode(codes)
|
||||
return quantized
|
||||
|
||||
class HiggsAudioTokenizer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
D: int = 256,
|
||||
target_bandwidths= [0.5, 1, 1.5, 2, 4],
|
||||
ratios = [8, 5, 4, 2, 3], # downsampling by 320
|
||||
sample_rate: int = 24000,
|
||||
bins: int = 1024,
|
||||
n_q: int = 8,
|
||||
codebook_dim: int = 64,
|
||||
last_layer_semantic: bool = True,
|
||||
downsample_mode: str = "step_down",
|
||||
vq_scale: int = 1,
|
||||
semantic_sample_rate: int = None,
|
||||
device = None,
|
||||
dtype = None,
|
||||
operations = None,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
operations = operations or nn
|
||||
self.hop_length = np.prod(ratios)
|
||||
|
||||
self.frame_rate = math.ceil(sample_rate / np.prod(ratios)) # 50 Hz
|
||||
|
||||
self.target_bandwidths = target_bandwidths
|
||||
self.n_q = n_q
|
||||
self.sample_rate = sample_rate
|
||||
self.encoder = DACEncoder(64, ratios, D, device = device, dtype = dtype, operations = operations)
|
||||
|
||||
self.decoder_2 = DACDecoder(D, 1024, ratios, device = device, dtype = dtype, operations = operations)
|
||||
self.last_layer_semantic = last_layer_semantic
|
||||
self.device = device
|
||||
|
||||
self.semantic_model = AutoModel.from_pretrained("bosonai/hubert_base", trust_remote_code=True)
|
||||
self.semantic_sample_rate = 16000
|
||||
self.semantic_dim = 768
|
||||
self.encoder_semantic_dim = 768
|
||||
|
||||
# Overwrite semantic model sr to ensure semantic_downsample_factor is an integer
|
||||
if semantic_sample_rate is not None:
|
||||
self.semantic_sample_rate = semantic_sample_rate
|
||||
|
||||
self.semantic_model.eval()
|
||||
|
||||
# make the semantic model parameters do not need gradient
|
||||
for param in self.semantic_model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.semantic_downsample_factor = int(self.hop_length / (self.sample_rate / self.semantic_sample_rate) / 320)
|
||||
|
||||
self.quantizer_dim = int((D + self.encoder_semantic_dim) // vq_scale)
|
||||
self.encoder_semantic = Encoder(input_channels=self.semantic_dim, encode_channels=self.encoder_semantic_dim, device = device, dtype = dtype, operations = operations)
|
||||
self.decoder_semantic = Decoder(
|
||||
code_dim=self.encoder_semantic_dim, output_channels=self.semantic_dim, decode_channels=self.semantic_dim, device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
self.quantizer = ResidualVectorQuantizer(
|
||||
dimension=self.quantizer_dim, codebook_dim=codebook_dim, n_q=n_q, bins=bins, device = device, dtype = dtype, operations = operations
|
||||
)
|
||||
|
||||
self.fc_prior = operations.Linear(D + self.encoder_semantic_dim, self.quantizer_dim, device = device, dtype = dtype)
|
||||
self.fc_post1 = operations.Linear(self.quantizer_dim, self.encoder_semantic_dim, device = device, dtype = dtype)
|
||||
self.fc_post2 = operations.Linear(self.quantizer_dim, D, device = device, dtype = dtype)
|
||||
|
||||
self.downsample_mode = downsample_mode
|
||||
|
||||
self.audio_tokenizer_feature_extractor = HiggsAudioFeatureExtractor(sampling_rate=self.sample_rate)
|
||||
|
||||
@property
|
||||
def sampling_rate(self):
|
||||
return self.sample_rate
|
||||
|
||||
@torch.no_grad()
|
||||
def get_regress_target(self, x):
|
||||
x = torchaudio.functional.resample(x, self.sample_rate, self.semantic_sample_rate)
|
||||
|
||||
x = x[:, 0, :]
|
||||
x = F.pad(x, (160, 160))
|
||||
target = self.semantic_model(x, output_hidden_states=True).hidden_states
|
||||
target = torch.stack(target, dim=1)
|
||||
|
||||
target = target.mean(1)
|
||||
|
||||
if self.downsample_mode == "step_down":
|
||||
if self.semantic_downsample_factor > 1:
|
||||
target = target[:, :: self.semantic_downsample_factor, :]
|
||||
|
||||
return target
|
||||
|
||||
def forward(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def tps(self):
|
||||
return self.frame_rate
|
||||
|
||||
def encode(self, wv, sr):
|
||||
|
||||
if sr != self.sampling_rate:
|
||||
# best computed values to match librosa's resample
|
||||
resampler_torch = torchaudio.transforms.Resample(
|
||||
orig_freq=sr,
|
||||
new_freq=self.sampling_rate,
|
||||
resampling_method="sinc_interp_kaiser",
|
||||
lowpass_filter_width = 121,
|
||||
rolloff = 0.9568384289091556,
|
||||
beta = 21.01531462440614
|
||||
).to(wv.device)
|
||||
|
||||
wv = resampler_torch(wv)
|
||||
|
||||
if self.audio_tokenizer_feature_extractor is not None:
|
||||
inputs = self.audio_tokenizer_feature_extractor(wv)
|
||||
input_values = inputs["input_values"].to(self.device)
|
||||
else:
|
||||
input_values = torch.from_numpy(wv).float().unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
input_values = input_values.to(wv.device)
|
||||
encoder_outputs = self._xcodec_encode(input_values)
|
||||
vq_code = encoder_outputs[0]
|
||||
return vq_code
|
||||
|
||||
def _xcodec_encode(self, x: torch.Tensor, target_bw: Optional[int] = None) -> torch.Tensor:
|
||||
bw = target_bw
|
||||
|
||||
e_semantic_input = self.get_regress_target(x).detach()
|
||||
|
||||
e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))
|
||||
e_acoustic = self.encoder(x)
|
||||
|
||||
if e_acoustic.shape[2] != e_semantic.shape[2]:
|
||||
pad_size = 160 * self.semantic_downsample_factor
|
||||
e_acoustic = self.encoder(F.pad(x[:, 0, :], (pad_size, pad_size)).unsqueeze(0))
|
||||
|
||||
if e_acoustic.shape[2] != e_semantic.shape[2]:
|
||||
if e_acoustic.shape[2] > e_semantic.shape[2]:
|
||||
e_acoustic = e_acoustic[:, :, : e_semantic.shape[2]]
|
||||
else:
|
||||
e_semantic = e_semantic[:, :, : e_acoustic.shape[2]]
|
||||
|
||||
e = torch.cat([e_acoustic, e_semantic], dim=1)
|
||||
|
||||
e = self.fc_prior(e.transpose(1, 2))
|
||||
|
||||
e = e.transpose(1, 2)
|
||||
_, codes, _, _ = self.quantizer(e, self.frame_rate, bw)
|
||||
codes = codes.permute(1, 0, 2)
|
||||
|
||||
return codes
|
||||
|
||||
def decode(self, vq_code: torch.Tensor) -> torch.Tensor:
|
||||
vq_code = vq_code.to(self.device)
|
||||
|
||||
if vq_code.ndim < 3:
|
||||
vq_code = vq_code.unsqueeze(0)
|
||||
|
||||
vq_code = vq_code.permute(1, 0, 2)
|
||||
quantized = self.quantizer.decode(vq_code)
|
||||
quantized = quantized.transpose(1, 2)
|
||||
quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)
|
||||
|
||||
o = self.decoder_2(quantized_acoustic)
|
||||
return o.detach()
|
||||
@ -46,6 +46,7 @@ import comfy.ldm.chroma.model
|
||||
import comfy.ldm.chroma_radiance.model
|
||||
import comfy.ldm.ace.model
|
||||
import comfy.ldm.omnigen.omnigen2
|
||||
import comfy.ldm.higgsv2.model
|
||||
import comfy.ldm.qwen_image.model
|
||||
|
||||
import comfy.model_management
|
||||
@ -1471,6 +1472,10 @@ class Omnigen2(BaseModel):
|
||||
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||
return out
|
||||
|
||||
class Higgsv2(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.EPS, device=None, unet_model=comfy.ldm.higgsv2.model.HiggsAudioModel):
|
||||
super().__init__(model_config, model_type, device, unet_model)
|
||||
|
||||
class QwenImage(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)
|
||||
|
||||
@ -527,6 +527,73 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
|
||||
return dit_config
|
||||
|
||||
if "{}layers.27.audio_post_attention_layernorm.weight".format(key_prefix) in state_dict_keys:
|
||||
|
||||
autoregressive_config = {}
|
||||
autoregressive_config["image_model"] = "higgsv2"
|
||||
|
||||
autoregressive_config["audio_adapter_type"] = "dual_ffn_fast_forward"
|
||||
autoregressive_config["audio_bos_token"] = "<|audio_bos|>"
|
||||
autoregressive_config["audio_codebook_size"] = 1024
|
||||
autoregressive_config["audio_num_codebooks"] = 8
|
||||
autoregressive_config["audio_ffn_hidden_size"] = 3072
|
||||
autoregressive_config["audio_ffn_intermediate_size"] = 8192
|
||||
autoregressive_config["audio_in_token"] = "<|AUDIO|>"
|
||||
autoregressive_config["audio_in_token_idx"] = 128015
|
||||
autoregressive_config["audio_out_token"] = "<|AUDIO_OUT|>"
|
||||
autoregressive_config["audio_out_token_idx"] = 128016
|
||||
autoregressive_config["audio_out_bos_token"] = "<|audio_out_bos|>"
|
||||
autoregressive_config["audio_out_bos_token_id"] = 128013
|
||||
autoregressive_config["audio_eos_token"] = "<|audio_eos|>"
|
||||
autoregressive_config["audio_eos_token_id"] = 128012
|
||||
autoregressive_config["audio_stream_bos_id"] = 1024
|
||||
autoregressive_config["audio_stream_eos_id"] = 1025
|
||||
autoregressive_config["encode_audio_in_tokens"] = True
|
||||
|
||||
autoregressive_config["pad_token_id"] = 128001
|
||||
autoregressive_config["padding_idx"] = 128001
|
||||
|
||||
autoregressive_config["hidden_size"] = 3072
|
||||
autoregressive_config["use_delay_pattern"] = True
|
||||
|
||||
autoregressive_config["vocab_size"] = 128256
|
||||
autoregressive_config["num_hidden_layers"] = 28
|
||||
autoregressive_config["num_attention_heads"] = 24
|
||||
autoregressive_config["num_key_value_heads"] = 8
|
||||
autoregressive_config["max_seq_len"] = 131072
|
||||
autoregressive_config["max_position_embeddings"] = 131072
|
||||
autoregressive_config["bos_token_id"] = 128000
|
||||
autoregressive_config["eos_token_id"] = 128001
|
||||
autoregressive_config["use_cache"] = True
|
||||
|
||||
autoregressive_config["text_config"] = {
|
||||
"model_type": "llama",
|
||||
"vocab_size": 128256,
|
||||
"max_position_embeddings": 131072,
|
||||
"num_hidden_layers": 28,
|
||||
"hidden_size": 3072,
|
||||
"num_attention_heads": 24,
|
||||
"num_key_value_heads": 8,
|
||||
"initializer_range": 0.02,
|
||||
"rms_norm_eps": 1e-05,
|
||||
"pad_token_id": None,
|
||||
"bos_token_id": 128000,
|
||||
"eos_token_id": 128001,
|
||||
"num_return_sequences": 1,
|
||||
"head_dim": 128,
|
||||
"mlp_bias": False,
|
||||
"intermediate_size": 8192
|
||||
}
|
||||
|
||||
autoregressive_config["use_kv_buckets"] = True
|
||||
autoregressive_config["num_attention_heads"] = 24
|
||||
|
||||
autoregressive_config["audio_decoder_proj_num_layers"] = 0
|
||||
autoregressive_config["audio_dual_ffn_layers"] = list(range(28))
|
||||
autoregressive_config["output_vae"] = False
|
||||
|
||||
return autoregressive_config
|
||||
|
||||
if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys: # HiDream
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "hidream"
|
||||
|
||||
@ -1312,6 +1312,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
|
||||
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
output_vae = model_config.unet_config.get("output_vae", output_vae)
|
||||
|
||||
if model_config.clip_vision_prefix is not None:
|
||||
if output_clipvision:
|
||||
|
||||
@ -19,6 +19,7 @@ import comfy.text_encoders.lumina2
|
||||
import comfy.text_encoders.wan
|
||||
import comfy.text_encoders.ace
|
||||
import comfy.text_encoders.omnigen2
|
||||
import comfy.text_encoders.higgsv2
|
||||
import comfy.text_encoders.qwen_image
|
||||
import comfy.text_encoders.hunyuan_image
|
||||
import comfy.text_encoders.z_image
|
||||
@ -1384,6 +1385,22 @@ class QwenImage(supported_models_base.BASE):
|
||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
||||
|
||||
class Higgsv2(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "higgsv2",
|
||||
}
|
||||
|
||||
memory_usage_factor = 1.0
|
||||
supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
text_encoder_key_prefix = ["dac."]
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
out = model_base.Higgsv2(self, device=device)
|
||||
return out
|
||||
|
||||
def clip_target(self, state_dict = {}):
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.higgsv2.DummyTokenizer, comfy.text_encoders.higgsv2.HiggsTokenizer)
|
||||
|
||||
class HunyuanImage21(HunyuanVideo):
|
||||
unet_config = {
|
||||
"image_model": "hunyuan_video",
|
||||
@ -1474,7 +1491,6 @@ class HunyuanVideo15_SR_Distilled(HunyuanVideo):
|
||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_video.HunyuanVideo15Tokenizer, comfy.text_encoders.hunyuan_image.te(**hunyuan_detect))
|
||||
|
||||
models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage, Flux2]
|
||||
|
||||
models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage, Flux2, Higgsv2]
|
||||
|
||||
models += [SVD_img2vid]
|
||||
|
||||
@ -0,0 +1,16 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|begin_of_text|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|end_of_text|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
1251004
comfy/text_encoders/higgs_text_tokenizer/tokenizer.json
Normal file
1251004
comfy/text_encoders/higgs_text_tokenizer/tokenizer.json
Normal file
File diff suppressed because it is too large
Load Diff
2063
comfy/text_encoders/higgs_text_tokenizer/tokenizer_config.json
Normal file
2063
comfy/text_encoders/higgs_text_tokenizer/tokenizer_config.json
Normal file
File diff suppressed because it is too large
Load Diff
81
comfy/text_encoders/higgsv2.py
Normal file
81
comfy/text_encoders/higgsv2.py
Normal file
@ -0,0 +1,81 @@
|
||||
import os
|
||||
import torch
|
||||
import comfy.ops
|
||||
import torch.nn as nn
|
||||
from transformers import AutoTokenizer
|
||||
from comfy.ldm.higgsv2.tokenizer import HiggsAudioTokenizer
|
||||
from comfy.ldm.higgsv2.preprocess import HiggsAudioSampleCollator
|
||||
|
||||
class DummyTokenizer:
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
pass
|
||||
|
||||
def revert_delay_pattern_vectorized(data: torch.Tensor) -> torch.Tensor:
|
||||
num_codebooks, total_len = data.shape
|
||||
seq_len = total_len - num_codebooks + 1
|
||||
|
||||
col_idx = torch.arange(seq_len, device=data.device)[None, :] \
|
||||
+ torch.arange(num_codebooks, device=data.device)[:, None]
|
||||
out = data[torch.arange(num_codebooks)[:, None], col_idx]
|
||||
return out
|
||||
|
||||
class HiggsTokenizer(nn.Module):
|
||||
def __init__(self, device, dtype, model_options={}, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
self.dtype = torch.float32
|
||||
self.device = device
|
||||
self.dtypes = [torch.float32]
|
||||
|
||||
here = os.path.dirname(__file__)
|
||||
tokenizer_path = os.path.join(here, "higgs_text_tokenizer")
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
scaled_fp8 = model_options.get("scaled_fp8", None)
|
||||
|
||||
if scaled_fp8 is not None:
|
||||
operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8)
|
||||
else:
|
||||
operations = comfy.ops.manual_cast
|
||||
|
||||
self.audio_codebook_size = 1024
|
||||
self.audio_tokenizer = HiggsAudioTokenizer(device = device, dtype = dtype, operations = operations)
|
||||
|
||||
if scaled_fp8 is not None:
|
||||
self.audio_tokenizer.scaled_fp8 = torch.nn.Parameter(torch.tensor([], dtype=scaled_fp8))
|
||||
|
||||
self.collator = HiggsAudioSampleCollator(
|
||||
audio_in_token_id = 128015,
|
||||
audio_out_token_id = 128016,
|
||||
audio_stream_bos_id = 1024,
|
||||
audio_stream_eos_id = 1025,
|
||||
pad_token_id = 128001,
|
||||
return_audio_in_tokens = False,
|
||||
use_delay_pattern = True,
|
||||
audio_num_codebooks = 8,
|
||||
round_to = 1,
|
||||
)
|
||||
|
||||
postfix = "<|start_header_id|>assistant<|end_header_id|>\n\n"
|
||||
self.postfix = postfix + "<|audio_out_bos|>" # force audio generation
|
||||
|
||||
def decode_tokens(self, audio_tokens):
|
||||
outputs = []
|
||||
|
||||
# due to instability issues, I had to convert the audio tokenizer to float32, avoiding outputing nans
|
||||
self.audio_tokenizer = self.audio_tokenizer.to(self.dtype)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for audio in audio_tokens:
|
||||
vq_code = revert_delay_pattern_vectorized(audio).clip(0, self.audio_codebook_size - 1)[:, 1:-1]
|
||||
wv_numpy = self.audio_tokenizer.decode(vq_code.unsqueeze(0))[0, 0]
|
||||
outputs.append(wv_numpy)
|
||||
|
||||
# currently only supports one batch size
|
||||
return (None, {"waveform": torch.cat(outputs, dim = 0).unsqueeze(0).unsqueeze(1), "sample_rate": self.audio_tokenizer.sample_rate}) # audio only
|
||||
|
||||
def load_state_dict(self, sd, strict = False):
|
||||
return self.audio_tokenizer.load_state_dict(sd, strict = strict)
|
||||
|
||||
def state_dict(self):
|
||||
return self.audio_tokenizer.state_dict()
|
||||
@ -1,12 +1,13 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Any
|
||||
from dataclasses import dataclass, field
|
||||
from transformers.cache_utils import Cache
|
||||
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
|
||||
import math
|
||||
import logging
|
||||
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
import comfy.model_management
|
||||
import comfy.ldm.common_dit
|
||||
|
||||
import comfy.model_management
|
||||
@ -24,9 +25,20 @@ class Llama2Config:
|
||||
rms_norm_eps: float = 1e-5
|
||||
rope_theta: float = 500000.0
|
||||
transformer_type: str = "llama"
|
||||
head_dim = 128
|
||||
head_dim: int = 128
|
||||
rms_norm_add = False
|
||||
mlp_activation = "silu"
|
||||
qkv_bias: bool = False
|
||||
rope_type: str = "llama3"
|
||||
rope_scaling: dict = field(
|
||||
default_factory=lambda: {
|
||||
"factor": 32.0,
|
||||
"high_freq_factor": 4.0,
|
||||
"low_freq_factor": 1.0,
|
||||
"original_max_position_embeddings": 8192,
|
||||
"rope_type": "llama3"
|
||||
}
|
||||
)
|
||||
qkv_bias = False
|
||||
rope_dims = None
|
||||
q_norm = None
|
||||
@ -255,15 +267,67 @@ def apply_rope(xq, xk, freqs_cis):
|
||||
sin = freqs_cis[1]
|
||||
q_embed = (xq * cos) + (rotate_half(xq) * sin)
|
||||
k_embed = (xk * cos) + (rotate_half(xk) * sin)
|
||||
return q_embed.to(org_dtype), k_embed.to(org_dtype)
|
||||
return q_embed.to(org_dtype), k_embed.to(org_dtype), sin, cos
|
||||
|
||||
class LlamaRoPE(nn.Module):
|
||||
def __init__(self, config, device = None, dtype = None):
|
||||
super().__init__()
|
||||
|
||||
if config.rope_scaling is not None:
|
||||
self.rope_type = config.rope_scaling.get("rope_type", config.rope_type)
|
||||
else:
|
||||
self.rope_type = "default"
|
||||
|
||||
self.config = config
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.original_inv_freq = self.inv_freq
|
||||
|
||||
def _dynamic_frequency_update(self, position_ids, device):
|
||||
|
||||
seq_len = torch.max(position_ids) + 1
|
||||
if seq_len > self.max_seq_len_cached:
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(
|
||||
self.config, device, seq_len=seq_len, **self.rope_kwargs
|
||||
)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.max_seq_len_cached = seq_len
|
||||
|
||||
if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:
|
||||
self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
|
||||
self.max_seq_len_cached = self.original_max_seq_len
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, x, position_ids):
|
||||
if "dynamic" in self.rope_type:
|
||||
self._dynamic_frequency_update(position_ids, device=x.device)
|
||||
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
|
||||
device_type = x.device.type
|
||||
device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
|
||||
with torch.autocast(device_type=device_type, enabled=False):
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos = emb.cos()
|
||||
sin = emb.sin()
|
||||
|
||||
cos = cos * self.attention_scaling
|
||||
sin = sin * self.attention_scaling
|
||||
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
|
||||
def __init__(self, config: Llama2Config, layer_idx: int = None, device=None, dtype=None, ops: Any = None):
|
||||
super().__init__()
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
self.hidden_size = config.hidden_size
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
self.head_dim = config.head_dim
|
||||
self.inner_size = self.num_heads * self.head_dim
|
||||
@ -286,6 +350,8 @@ class Attention(nn.Module):
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_value: Optional[Cache] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
freqs_cis: Optional[torch.Tensor] = None,
|
||||
optimized_attention=None,
|
||||
):
|
||||
@ -303,13 +369,22 @@ class Attention(nn.Module):
|
||||
if self.k_norm is not None:
|
||||
xk = self.k_norm(xk)
|
||||
|
||||
xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
|
||||
xq, xk, sin, cos = apply_rope(xq, xk, freqs_cis=freqs_cis)
|
||||
|
||||
if past_key_value is not None:
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
||||
xk, xv = past_key_value.update(xk, xv, self.layer_idx, cache_kwargs)
|
||||
|
||||
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
|
||||
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
|
||||
return self.o_proj(output)
|
||||
out = self.o_proj(output)
|
||||
|
||||
if past_key_value is not None:
|
||||
return out, past_key_value
|
||||
|
||||
return out
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import av
|
||||
import re
|
||||
import torchaudio
|
||||
import torch
|
||||
import comfy.model_management
|
||||
@ -10,10 +11,258 @@ import io
|
||||
import json
|
||||
import random
|
||||
import hashlib
|
||||
import numpy as np
|
||||
import node_helpers
|
||||
import logging
|
||||
from comfy.cli_args import args
|
||||
from comfy.comfy_types import IO
|
||||
from comfy.comfy_types import FileLocator
|
||||
from dataclasses import asdict
|
||||
from comfy.ldm.higgsv2.loudness import loudness
|
||||
from comfy.ldm.higgsv2.preprocess import (
|
||||
prepare_chatml_sample, Message, ChatMLSample, ChatMLDatasetSample, AudioContent, transcript_normalize
|
||||
)
|
||||
|
||||
|
||||
MULTISPEAKER_DEFAULT_SYSTEM_MESSAGE = """You are an AI assistant designed to convert text into speech.
|
||||
If the user's message includes a [SPEAKER*] tag, do not read out the tag and generate speech for the following text, using the specified voice.
|
||||
If no speaker tag is present, select a suitable voice on your own."""
|
||||
|
||||
class LoudnessNormalization:
|
||||
|
||||
CATEGORY = "audio"
|
||||
RETURN_TYPES = ("AUDIO",)
|
||||
FUNCTION = "normalize"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"audio": ("AUDIO", ),
|
||||
"block_size": ("FLOAT", {"default": 0.400, "min": 0.1, "max": 1.0, "step": 0.05}),
|
||||
"loudness_threshold": ("FLOAT", {"default": -23.0, "min": -70.0, "max": 0.0, "step": 0.5,
|
||||
"tooltip": "Target loudness in LUFS. Common values are -23.0 (broadcast), -14.0 (streaming)."})}}
|
||||
|
||||
def normalize(self, audio, loudness_threshold, block_size):
|
||||
sampling_rate = audio["sample_rate"]
|
||||
waveform = audio["waveform"]
|
||||
return {"waveform": loudness(waveform, sampling_rate, target_loudness = loudness_threshold, block_size = block_size), "sample_rate": sampling_rate}
|
||||
|
||||
def prepare_chatml_input(
|
||||
clip,
|
||||
input_tokens,
|
||||
audio_contents,
|
||||
sampling_rate,
|
||||
postfix_str: str = "",
|
||||
):
|
||||
if hasattr(clip, "postfix"):
|
||||
postfix_str = clip.postfix
|
||||
|
||||
if postfix_str:
|
||||
postfix = clip.tokenizer.encode(postfix_str, add_special_tokens=False)
|
||||
input_tokens.extend(postfix)
|
||||
|
||||
audio_ids_l = []
|
||||
if audio_contents is not None:
|
||||
|
||||
if not hasattr(clip, "audio_tokenizer"):
|
||||
raise ValueError("This model does not have an audio tokenizer. The chosen model may not support ChatML Format")
|
||||
|
||||
for audio_content in audio_contents:
|
||||
audio_content.raw_audio = audio_content.raw_audio.squeeze(0)
|
||||
if audio_content.raw_audio.shape[0] == 2:
|
||||
audio_content.raw_audio = audio_content.raw_audio.mean(dim = 0, keepdim = True)
|
||||
|
||||
if audio_content.raw_audio.device != next(clip.audio_tokenizer.parameters()).device:
|
||||
audio_content.raw_audio = audio_content.raw_audio.to(next(clip.audio_tokenizer.parameters()).device)
|
||||
|
||||
audio_ids = clip.audio_tokenizer.encode(audio_content.raw_audio, sampling_rate)
|
||||
audio_ids_l.append(audio_ids.squeeze(0))
|
||||
|
||||
if len(audio_ids_l) > 0:
|
||||
audio_ids_start = torch.tensor(
|
||||
np.cumsum(np.array([0] + [audio_ids.shape[1] for audio_ids in audio_ids_l])),
|
||||
dtype=torch.long,
|
||||
device=audio_contents[0].raw_audio.device,
|
||||
).to("cpu")[0:-1]
|
||||
audio_ids_concat = torch.cat(audio_ids_l, dim=1).to("cpu")
|
||||
else:
|
||||
audio_ids_start = None
|
||||
audio_ids_concat = None
|
||||
|
||||
sample = ChatMLDatasetSample(
|
||||
input_ids=torch.LongTensor(input_tokens),
|
||||
label_ids=None,
|
||||
audio_ids_concat=audio_ids_concat,
|
||||
audio_ids_start=audio_ids_start,
|
||||
audio_waveforms_concat=None,
|
||||
audio_waveforms_start=None,
|
||||
audio_sample_rate=None,
|
||||
audio_speaker_indices=None,
|
||||
)
|
||||
|
||||
if hasattr(clip, "collator"):
|
||||
sample.input_ids = sample.input_ids.cpu()
|
||||
sample = clip.collator([sample])
|
||||
|
||||
inputs = asdict(sample)
|
||||
for k, v in inputs.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
inputs[k] = v.to(clip.device)
|
||||
|
||||
return inputs
|
||||
|
||||
def postprocess_chatml(text: str) -> str:
|
||||
speakers = set(re.findall(r'\[SPEAKER\d+\]', text))
|
||||
skip_recon = True
|
||||
|
||||
if len(speakers) > 1:
|
||||
parts = text.split('<|eot_id|>')
|
||||
|
||||
# keep the first <|eot_id|> and the last one
|
||||
first_eot = parts[0] + '<|eot_id|>'
|
||||
middle_parts = ''.join(parts[1:-1])
|
||||
last_eot = '<|eot_id|>' + parts[-1]
|
||||
|
||||
text = first_eot + middle_parts + last_eot
|
||||
skip_recon = False
|
||||
|
||||
return text, skip_recon
|
||||
|
||||
class CreateChatMLSample:
|
||||
def __init__(self):
|
||||
self.device = comfy.model_management.intermediate_device()
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": (IO.STRING, {
|
||||
"default": "SYSTEM: " + MULTISPEAKER_DEFAULT_SYSTEM_MESSAGE + "\n\n<|scene_desc_start|>\nSPEAKER0:masculine\nSPEAKER1:feminine\n<|scene_desc_end|>",
|
||||
"multiline": True,
|
||||
"dynamicPrompts": True,
|
||||
"tooltip": (
|
||||
"The conversations to be encoded. "
|
||||
"To register a conversation start with SPEAKER-0: some text. "
|
||||
"To add a system prompt start with system:"
|
||||
),
|
||||
}),
|
||||
"clip": (IO.CLIP, {"tooltip": "The CLIP model used for tokenizing the text."}),
|
||||
},
|
||||
"optional": {
|
||||
"audio": (IO.AUDIO, {
|
||||
"tooltip": "An audio clip to be inserted into the conversation. To register add [audio]",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("TOKENS",)
|
||||
OUTPUT_TOOLTIPS = ("Turns text and audio into a ChatML Format.",)
|
||||
|
||||
FUNCTION = "convert_to_ml_format"
|
||||
CATEGORY = "conditioning"
|
||||
|
||||
def convert_to_ml_format(self, clip, text, audio=None):
|
||||
|
||||
if audio is not None:
|
||||
clip.load_model()
|
||||
|
||||
if hasattr(clip, "cond_stage_model"):
|
||||
clip = clip.cond_stage_model
|
||||
|
||||
text = transcript_normalize(text)
|
||||
|
||||
messages = []
|
||||
lines = text.splitlines()
|
||||
sampling_rate = False
|
||||
current_role = None
|
||||
collecting_system = False
|
||||
system_buffer = []
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# system start
|
||||
if line.lower().startswith("system:"):
|
||||
collecting_system = True
|
||||
system_buffer.append(line[len("system:"):].strip())
|
||||
continue
|
||||
|
||||
# while collecting system prompt
|
||||
if collecting_system:
|
||||
system_buffer.append(line)
|
||||
if "<|scene_desc_end|>" in line or "SPEAKER-" in line:
|
||||
system_prompt = "\n".join(system_buffer)# + "\n<|scene_desc_end|>"
|
||||
messages.append(Message(role="system", content=system_prompt))
|
||||
system_buffer = []
|
||||
collecting_system = False
|
||||
continue
|
||||
|
||||
# speaker lines SPEAKER-0: text
|
||||
match = re.match(r"SPEAKER-(\d+):\s*(.*)", line, re.IGNORECASE)
|
||||
if match:
|
||||
speaker_id = match.group(1)
|
||||
content = match.group(2)
|
||||
current_role = f"[SPEAKER{speaker_id}] "
|
||||
messages.append(Message(role = "user", content = current_role + content.strip()))
|
||||
else:
|
||||
# continuation line goes to last speaker or instruction
|
||||
if current_role is not None and messages:
|
||||
messages[-1].content += "\n" + line
|
||||
|
||||
# return normal input_ids
|
||||
if not (len(messages) >= 1):
|
||||
return (clip.tokenizer(text),)
|
||||
|
||||
all_text = "".join(msg.content for msg in messages if msg.role == "user")
|
||||
|
||||
# postprocess to allow multi-user speech
|
||||
all_text, skip_recon = postprocess_chatml(all_text)
|
||||
if not skip_recon:
|
||||
lines = all_text.splitlines()
|
||||
messages = [messages[0]] if messages[0].role == "system" else []
|
||||
current_role = None
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
match = re.match(r'(\[SPEAKER\d+\])\s*(.*)', line)
|
||||
if match:
|
||||
current_role = match.group(1)
|
||||
content = match.group(2).strip() # only take the text after the tag
|
||||
messages.append(Message(role="user", content=f"{current_role} {content}" if content else current_role))
|
||||
else:
|
||||
if current_role and messages:
|
||||
messages[-1].content += "\n" + line
|
||||
|
||||
# dedepulicate the messages
|
||||
for idx, m in enumerate(messages):
|
||||
double_eot = "<|eot_id|><|eot_id|>"
|
||||
if double_eot in m.content:
|
||||
cut_index = m.content.index(double_eot)
|
||||
messages[idx].content = m.content[:cut_index + (len(double_eot) // 2)]
|
||||
break
|
||||
|
||||
if audio is not None:
|
||||
# for audio cloning, the first message is a transcript, second is the audio,
|
||||
# third is the request of what the model should say
|
||||
waveform = audio["waveform"]
|
||||
sampling_rate = audio["sample_rate"]
|
||||
messages.insert(1, Message(
|
||||
role = "assistant",
|
||||
content = AudioContent(raw_audio = waveform, audio_url = "placeholder")
|
||||
))
|
||||
chat_ml_sample = ChatMLSample(messages)
|
||||
input_tokens, audio_contents, _ = prepare_chatml_sample(
|
||||
chat_ml_sample,
|
||||
clip.tokenizer,
|
||||
)
|
||||
|
||||
if audio is None:
|
||||
audio_contents = None
|
||||
out = prepare_chatml_input(clip, input_tokens, audio_contents, sampling_rate = sampling_rate)
|
||||
return (out,)
|
||||
|
||||
class EmptyLatentAudio:
|
||||
def __init__(self):
|
||||
@ -586,6 +835,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LoadAudio": LoadAudio,
|
||||
"PreviewAudio": PreviewAudio,
|
||||
"ConditioningStableAudio": ConditioningStableAudio,
|
||||
"LoudnessNormalization": LoudnessNormalization,
|
||||
"CreateChatMLSample": CreateChatMLSample,
|
||||
"RecordAudio": RecordAudio,
|
||||
"TrimAudioDuration": TrimAudioDuration,
|
||||
"SplitAudioChannels": SplitAudioChannels,
|
||||
@ -604,6 +855,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SaveAudio": "Save Audio (FLAC)",
|
||||
"SaveAudioMP3": "Save Audio (MP3)",
|
||||
"SaveAudioOpus": "Save Audio (Opus)",
|
||||
"LoudnessNormalization": "Loudness Normalization",
|
||||
"CreateChatMLSample": "Create ChatML Sample",
|
||||
"RecordAudio": "Record Audio",
|
||||
"TrimAudioDuration": "Trim Audio Duration",
|
||||
"SplitAudioChannels": "Split Audio Channels",
|
||||
|
||||
60
comfy_extras/nodes_autoregressive.py
Normal file
60
comfy_extras/nodes_autoregressive.py
Normal file
@ -0,0 +1,60 @@
|
||||
from comfy.autoregressive_sampling import auto_sample
|
||||
from comfy.comfy_types import IO
|
||||
|
||||
class AutoRegressiveGeneration:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The model used for generation."}),
|
||||
"input_ids": ("TOKENS", ),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for controling the generation."}),
|
||||
"max_new_length": ("INT", {"default": 1024, "min": 1, "max": 10_000, "tooltip": "The max length for generation."}),
|
||||
"min_new_length": ("INT", {"default": 1, "min": 1, "max": 10_000, "tooltip": "The min length for generation."}),
|
||||
"top_k": ("INT", {"default": 50, "min": 1, "max": 30_000, "tooltip": "Takes the top k of the most probable tokens."}),
|
||||
"top_p": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Percentage of tokens to leave after generation (top most probable tokens)."}),
|
||||
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 50, "step": 0.01, "tooltip": "Temperature controls randomess by decreasing or increasing the probability of lesser likely tokens. Higher Temperature -> More Randomness"}),
|
||||
"do_sample": ("BOOLEAN", {"default": False, "tooltip": "Add randomness in decoding the tokens."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GENERATED_TOKENS",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
CATEGORY = "sampling"
|
||||
|
||||
# for cuda graphs
|
||||
_cached_autoregressive_sampler = None
|
||||
|
||||
def generate(self, model, input_ids, seed, max_new_length, min_new_length, top_k, top_p, temperature, do_sample):
|
||||
return (auto_sample(self, model, input_ids, max_new_length, min_new_length, top_k, top_p, temperature, do_sample, seed = seed),)
|
||||
|
||||
class DecodeTokens:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"clip": (IO.CLIP, {"tooltip": "The model used for generation."}),
|
||||
"generated_tokens": ("GENERATED_TOKENS", ),}
|
||||
}
|
||||
|
||||
FUNCTION = "decode"
|
||||
CATEGORY = "conditioning"
|
||||
RETURN_TYPES = ("TEXT", "AUDIO")
|
||||
|
||||
def decode(self, clip, tokens):
|
||||
clip.load_model()
|
||||
if hasattr(clip.cond_stage_model, "decode_tokens"): # for special tokenizers
|
||||
return clip.cond_stage_model.decode_tokens(tokens)
|
||||
else:
|
||||
return (clip.tokenizer.decode(tokens, skip_special_tokens=True), None)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DecodeTokens": DecodeTokens,
|
||||
"AutoRegressiveGeneration": AutoRegressiveGeneration,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AutoRegressiveGeneration": "Autoregressive Generation",
|
||||
"DecodeTokens": "Decode Tokens",
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user