mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-23 10:03:36 +08:00
Refactor Helios to reuse WAN text encoder, latent format, and VAE
This commit is contained in:
parent
f9d26fc23f
commit
a5c328871d
@ -783,11 +783,3 @@ class ZImagePixelSpace(ChromaRadiance):
|
|||||||
No VAE encoding/decoding — the model operates directly on RGB pixels.
|
No VAE encoding/decoding — the model operates directly on RGB pixels.
|
||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
class Helios(Wan21):
|
|
||||||
"""Helios video model latent format
|
|
||||||
|
|
||||||
Helios uses the same latent format as Wan21 (same VAE architecture).
|
|
||||||
Inherits latents_mean, latents_std, and processing methods from Wan21.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|||||||
@ -48,7 +48,6 @@ import comfy.text_encoders.hunyuan_video
|
|||||||
import comfy.text_encoders.cosmos
|
import comfy.text_encoders.cosmos
|
||||||
import comfy.text_encoders.lumina2
|
import comfy.text_encoders.lumina2
|
||||||
import comfy.text_encoders.wan
|
import comfy.text_encoders.wan
|
||||||
import comfy.text_encoders.helios
|
|
||||||
import comfy.text_encoders.hidream
|
import comfy.text_encoders.hidream
|
||||||
import comfy.text_encoders.ace
|
import comfy.text_encoders.ace
|
||||||
import comfy.text_encoders.omnigen2
|
import comfy.text_encoders.omnigen2
|
||||||
|
|||||||
@ -17,7 +17,6 @@ import comfy.text_encoders.hunyuan_video
|
|||||||
import comfy.text_encoders.cosmos
|
import comfy.text_encoders.cosmos
|
||||||
import comfy.text_encoders.lumina2
|
import comfy.text_encoders.lumina2
|
||||||
import comfy.text_encoders.wan
|
import comfy.text_encoders.wan
|
||||||
import comfy.text_encoders.helios
|
|
||||||
import comfy.text_encoders.ace
|
import comfy.text_encoders.ace
|
||||||
import comfy.text_encoders.omnigen2
|
import comfy.text_encoders.omnigen2
|
||||||
import comfy.text_encoders.qwen_image
|
import comfy.text_encoders.qwen_image
|
||||||
@ -1143,7 +1142,7 @@ class Helios(supported_models_base.BASE):
|
|||||||
}
|
}
|
||||||
|
|
||||||
unet_extra_config = {}
|
unet_extra_config = {}
|
||||||
latent_format = latent_formats.Helios
|
latent_format = latent_formats.Wan21
|
||||||
memory_usage_factor = 1.8
|
memory_usage_factor = 1.8
|
||||||
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
||||||
|
|
||||||
@ -1159,8 +1158,15 @@ class Helios(supported_models_base.BASE):
|
|||||||
|
|
||||||
def clip_target(self, state_dict={}):
|
def clip_target(self, state_dict={}):
|
||||||
pref = self.text_encoder_key_prefix[0]
|
pref = self.text_encoder_key_prefix[0]
|
||||||
t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}umt5xxl.transformer.".format(pref))
|
t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(
|
||||||
return supported_models_base.ClipTarget(comfy.text_encoders.helios.HeliosT5Tokenizer, comfy.text_encoders.helios.te(**t5_detect))
|
state_dict,
|
||||||
|
"{}umt5xxl.transformer.".format(pref),
|
||||||
|
)
|
||||||
|
# Directly reuse WAN text encoder stack; no Helios-specific TE.
|
||||||
|
return supported_models_base.ClipTarget(
|
||||||
|
comfy.text_encoders.wan.WanT5Tokenizer,
|
||||||
|
comfy.text_encoders.wan.te(**t5_detect),
|
||||||
|
)
|
||||||
|
|
||||||
class WAN21_T2V(supported_models_base.BASE):
|
class WAN21_T2V(supported_models_base.BASE):
|
||||||
unet_config = {
|
unet_config = {
|
||||||
|
|||||||
@ -1,41 +0,0 @@
|
|||||||
from comfy import sd1_clip
|
|
||||||
from .spiece_tokenizer import SPieceTokenizer
|
|
||||||
import comfy.text_encoders.t5
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
class UMT5XXlModel(sd1_clip.SDClipModel):
|
|
||||||
def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}):
|
|
||||||
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "umt5_config_xxl.json")
|
|
||||||
super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=True, zero_out_masked=True, model_options=model_options)
|
|
||||||
|
|
||||||
|
|
||||||
class UMT5XXlTokenizer(sd1_clip.SDTokenizer):
|
|
||||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
||||||
tokenizer = tokenizer_data.get("spiece_model", None)
|
|
||||||
super().__init__(tokenizer, pad_with_end=False, embedding_size=4096, embedding_key="umt5xxl", tokenizer_class=SPieceTokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_token=0, tokenizer_data=tokenizer_data)
|
|
||||||
|
|
||||||
def state_dict(self):
|
|
||||||
return {"spiece_model": self.tokenizer.serialize_model()}
|
|
||||||
|
|
||||||
|
|
||||||
class HeliosT5Tokenizer(sd1_clip.SD1Tokenizer):
|
|
||||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
||||||
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="umt5xxl", tokenizer=UMT5XXlTokenizer)
|
|
||||||
|
|
||||||
|
|
||||||
class HeliosT5Model(sd1_clip.SD1ClipModel):
|
|
||||||
def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs):
|
|
||||||
super().__init__(device=device, dtype=dtype, model_options=model_options, name="umt5xxl", clip_model=UMT5XXlModel, **kwargs)
|
|
||||||
|
|
||||||
|
|
||||||
def te(dtype_t5=None, t5_quantization_metadata=None):
|
|
||||||
class HeliosTEModel(HeliosT5Model):
|
|
||||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
||||||
if t5_quantization_metadata is not None:
|
|
||||||
model_options = model_options.copy()
|
|
||||||
model_options["quantization_metadata"] = t5_quantization_metadata
|
|
||||||
if dtype_t5 is not None:
|
|
||||||
dtype = dtype_t5
|
|
||||||
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
|
||||||
return HeliosTEModel
|
|
||||||
@ -42,7 +42,7 @@ def _parse_int_list(values, default):
|
|||||||
return out if len(out) > 0 else default
|
return out if len(out) > 0 else default
|
||||||
|
|
||||||
|
|
||||||
_HELIOS_LATENT_FORMAT = comfy.latent_formats.Helios()
|
_HELIOS_LATENT_FORMAT = comfy.latent_formats.Wan21()
|
||||||
|
|
||||||
|
|
||||||
def _apply_helios_latent_space_noise(latent, sigma, generator=None):
|
def _apply_helios_latent_space_noise(latent, sigma, generator=None):
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user