Re-coded T5 loader and WanVideo Text Encoder 2/24 complete

__init__.py

@@ -117,6 +117,7 @@ from .nodes import (
 
 from .wanvideo import (
     LoadWanVideoT5TextEncoder,
+    WanVideoTextEncode,
 )
 
 from .wrappers import (
@@ -253,7 +254,8 @@ pulid_nodes = {
 register_and_count(["PuLID_ComfyUI", "pulid_comfyui"], pulid_nodes)
 
 wanvideo_nodes = {
-    "LoadWanVideoT5TextEncoderMultiGPU": LoadWanVideoT5TextEncoder
+    "LoadWanVideoT5TextEncoderMultiGPU": LoadWanVideoT5TextEncoder,
+    "WanVideoTextEncodeMultiGPU": WanVideoTextEncode
 }
 register_and_count(["ComfyUI-WanVideoWrapper", "comfyui-wanvideowrapper"], wanvideo_nodes)
 

wanvideo.py

@@ -29,8 +29,8 @@ class LoadWanVideoT5TextEncoder:
             }
         }
 
-    RETURN_TYPES = ("WANTEXTENCODER", "STRING")
-    RETURN_NAMES = ("wan_t5_model", "device")
+    RETURN_TYPES = ("WANTEXTENCODER", "MULTIGPUDEVICE")
+    RETURN_NAMES = ("wan_t5_model", "load_device")
     FUNCTION = "loadmodel"
     CATEGORY = "multigpu/WanVideoWrapper"
     DESCRIPTION = "Loads Wan text_encoder model from 'ComfyUI/models/LLM'"
@@ -46,11 +46,63 @@ class LoadWanVideoT5TextEncoder:
         else:
             load_device = "main_device"
 
-        logger.info(f"[MultiGPU WanVideoWrapper] current_device set to: {device}")
-        logger.info(f"[MultiGPU WanVideoWrapper] load_device set to: {load_device}")
+        logger.info(f"[MultiGPU WanVideoWrapper][LoadWanVideoT5TextEncoder] current_device set to: {device}")
+        logger.info(f"[MultiGPU WanVideoWrapper][LoadWanVideoT5TextEncoder] load_device set to: {load_device}")
 
         original_loader = NODE_CLASS_MAPPINGS["LoadWanVideoT5TextEncoder"]()
         text_encoder = original_loader.loadmodel(model_name, precision, load_device, quantization)
 
         # Return both the text encoder AND the selected device
         return text_encoder, device
+
+
+class WanVideoTextEncode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "positive_prompt": ("STRING", {"default": "", "multiline": True} ),
+            "negative_prompt": ("STRING", {"default": "", "multiline": True} ),
+            },
+            "optional": {
+                "t5": ("WANTEXTENCODER",),
+                "load_device": ("MULTIGPUDEVICE",),
+                "force_offload": ("BOOLEAN", {"default": True}),
+                "model_to_offload": ("WANVIDEOMODEL", {"tooltip": "Model to move to offload_device before encoding"}),
+                "use_disk_cache": ("BOOLEAN", {"default": False, "tooltip": "Cache the text embeddings to disk for faster re-use, under the custom_nodes/ComfyUI-WanVideoWrapper/text_embed_cache directory"}),
+                #"device": (["gpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}),
+            }
+        }
+
+    RETURN_TYPES = ("WANVIDEOTEXTEMBEDS", )
+    RETURN_NAMES = ("text_embeds",)
+    FUNCTION = "process"
+    CATEGORY = "multigpu/WanVideoWrapper"
+    DESCRIPTION = "Encodes text prompts into text embeddings. For rudimentary prompt travel you can input multiple prompts separated by '|', they will be equally spread over the video length"
+
+
+    def process(self, positive_prompt, negative_prompt, t5=None, load_device=None,force_offload=True, model_to_offload=None, use_disk_cache=False):
+        from . import set_current_device
+
+        if load_device is not None:
+            set_current_device(load_device)
+
+        if load_device == "cpu":
+            device = "cpu"
+        else:
+
+        if t5 is not None:
+            text_encoder = t5[0]
+        else:
+            text_encoder = None
+
+        logger.info(f"[MultiGPU WanVideoWrapper][WanVideoTextEncode] current_device set to: {load_device}")
+        logger.info(f"[MultiGPU WanVideoWrapper][WanVideoTextEncode] device set to: {device}")
+
+        original_encoder = NODE_CLASS_MAPPINGS["WanVideoTextEncode"]()
+        prompt_embeds_dict = original_encoder.process(positive_prompt, negative_prompt, text_encoder, force_offload, model_to_offload, use_disk_cache, device)
+        return (prompt_embeds_dict)
+
+    def parse_prompt_weights(self, prompt):
+        """Extract text and weights from prompts with (text:weight) format"""
+        original_parser = NODE_CLASS_MAPPINGS["WanVideoTextEncode"]()
+        return original_parser.parse_prompt_weights(prompt)