Refactor Helios to reuse WAN text encoder, latent format, and VAE

comfy/latent_formats.py

@@ -783,11 +783,3 @@ class ZImagePixelSpace(ChromaRadiance):
     No VAE encoding/decoding — the model operates directly on RGB pixels.
     """
     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

comfy/sd.py

@@ -48,7 +48,6 @@ import comfy.text_encoders.hunyuan_video
 import comfy.text_encoders.cosmos
 import comfy.text_encoders.lumina2
 import comfy.text_encoders.wan
-import comfy.text_encoders.helios
 import comfy.text_encoders.hidream
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2

comfy/supported_models.py

@@ -17,7 +17,6 @@ import comfy.text_encoders.hunyuan_video
 import comfy.text_encoders.cosmos
 import comfy.text_encoders.lumina2
 import comfy.text_encoders.wan
-import comfy.text_encoders.helios
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
@@ -1143,7 +1142,7 @@ class Helios(supported_models_base.BASE):
     }
 
     unet_extra_config = {}
-    latent_format = latent_formats.Helios
+    latent_format = latent_formats.Wan21
     memory_usage_factor = 1.8
     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={}):
         pref = self.text_encoder_key_prefix[0]
-        t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}umt5xxl.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.helios.HeliosT5Tokenizer, comfy.text_encoders.helios.te(**t5_detect))
+        t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_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):
     unet_config = {

comfy/text_encoders/helios.py

@@ -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

comfy_extras/nodes_helios.py

@@ -42,7 +42,7 @@ def _parse_int_list(values, 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):