Add WAN ATI support

ati.py

@@ -0,0 +1,223 @@
+import os
+import random
+import sys
+from typing import Sequence, Mapping, Any, Union
+import torch
+
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    """Returns the value at the given index of a sequence or mapping.
+
+    If the object is a sequence (like list or string), returns the value at the given index.
+    If the object is a mapping (like a dictionary), returns the value at the index-th key.
+
+    Some return a dictionary, in these cases, we look for the "results" key
+
+    Args:
+        obj (Union[Sequence, Mapping]): The object to retrieve the value from.
+        index (int): The index of the value to retrieve.
+
+    Returns:
+        Any: The value at the given index.
+
+    Raises:
+        IndexError: If the index is out of bounds for the object and the object is not a mapping.
+    """
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+
+def find_path(name: str, path: str = None) -> str:
+    """
+    Recursively looks at parent folders starting from the given path until it finds the given name.
+    Returns the path as a Path object if found, or None otherwise.
+    """
+    # If no path is given, use the current working directory
+    if path is None:
+        path = os.getcwd()
+
+    # Check if the current directory contains the name
+    if name in os.listdir(path):
+        path_name = os.path.join(path, name)
+        print(f"{name} found: {path_name}")
+        return path_name
+
+    # Get the parent directory
+    parent_directory = os.path.dirname(path)
+
+    # If the parent directory is the same as the current directory, we've reached the root and stop the search
+    if parent_directory == path:
+        return None
+
+    # Recursively call the function with the parent directory
+    return find_path(name, parent_directory)
+
+
+def add_comfyui_directory_to_sys_path() -> None:
+    """
+    Add 'ComfyUI' to the sys.path
+    """
+    comfyui_path = find_path("ComfyUI")
+    if comfyui_path is not None and os.path.isdir(comfyui_path):
+        sys.path.append(comfyui_path)
+        print(f"'{comfyui_path}' added to sys.path")
+
+
+def add_extra_model_paths() -> None:
+    """
+    Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
+    """
+    try:
+        from main import load_extra_path_config
+    except ImportError:
+        print(
+            "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
+        )
+        from utils.extra_config import load_extra_path_config
+
+    extra_model_paths = find_path("extra_model_paths.yaml")
+
+    if extra_model_paths is not None:
+        load_extra_path_config(extra_model_paths)
+    else:
+        print("Could not find the extra_model_paths config file.")
+
+
+add_comfyui_directory_to_sys_path()
+add_extra_model_paths()
+
+
+def import_custom_nodes() -> None:
+    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
+
+    This function sets up a new asyncio event loop, initializes the PromptServer,
+    creates a PromptQueue, and initializes the custom nodes.
+    """
+    import asyncio
+    import execution
+    from nodes import init_extra_nodes
+    import server
+
+    # Creating a new event loop and setting it as the default loop
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    # Creating an instance of PromptServer with the loop
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+
+    # Initializing custom nodes
+    init_extra_nodes()
+
+
+from nodes import NODE_CLASS_MAPPINGS
+
+
+def main():
+    import_custom_nodes()
+    with torch.inference_mode():
+        cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+        cliploader_38 = cliploader.load_clip(
+            clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors", type="wan", device="default"
+        )
+
+        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+        cliptextencode_6 = cliptextencode.encode(
+            text="A little boy and a dog are hiking in the forest. The puppy runs around the boy cheerfully.\n",
+            clip=get_value_at_index(cliploader_38, 0),
+        )
+
+        cliptextencode_7 = cliptextencode.encode(
+            text="过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，",
+            clip=get_value_at_index(cliploader_38, 0),
+        )
+
+        vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+        vaeloader_39 = vaeloader.load_vae(vae_name="kj-Wan2_1_VAE_fp32.safetensors")
+
+        unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+        unetloader_140 = unetloader.load_unet(
+            unet_name="Wan2_1-I2V-ATI-14B_fp8_e4m3fn.safetensors",
+            weight_dtype="fp8_e4m3fn_fast",
+        )
+
+        loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
+        loraloadermodelonly_154 = loraloadermodelonly.load_lora_model_only(
+            lora_name="Wan21_CausVid_14B_T2V_lora_rank32_v2.safetensors",
+            strength_model=0.30000000000000004,
+            model=get_value_at_index(unetloader_140, 0),
+        )
+
+        loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+        loadimage_240 = loadimage.load_image(image="example.png")
+
+        clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+        clipvisionloader_244 = clipvisionloader.load_clip(
+            clip_name="clip_vision_h.safetensors"
+        )
+
+        primitivestringmultiline = NODE_CLASS_MAPPINGS["PrimitiveStringMultiline"]()
+        primitivestringmultiline_247 = primitivestringmultiline.execute(
+            value="[[{'x':637.5706176757812,'y':357.4817810058594},{'x':637.1891479492188,'y':357.1156005859375},{'x':636.8077392578125,'y':356.7493896484375},{'x':636.42626953125,'y':356.3832092285156},{'x':636.0448608398438,'y':356.01702880859375},{'x':635.6633911132812,'y':355.6508483886719},{'x':635.2819213867188,'y':355.2846374511719},{'x':634.9005126953125,'y':354.91845703125},{'x':634.51904296875,'y':354.5522766113281},{'x':634.1375732421875,'y':354.18609619140625},{'x':633.7561645507812,'y':353.81988525390625},{'x':633.3746948242188,'y':353.4537048339844},{'x':632.9932861328125,'y':353.0875244140625},{'x':632.61181640625,'y':352.7213134765625},{'x':632.2305297851562,'y':352.35498046875},{'x':631.849365234375,'y':351.9884948730469},{'x':631.4681396484375,'y':351.6220397949219},{'x':631.0878295898438,'y':351.25469970703125},{'x':630.7080688476562,'y':350.8867492675781},{'x':630.328857421875,'y':350.51824951171875},{'x':629.9509887695312,'y':350.1484069824219},{'x':629.5737915039062,'y':349.7778015136719},{'x':629.197509765625,'y':349.4063415527344},{'x':628.8228759765625,'y':349.0331726074219},{'x':628.4490356445312,'y':348.65924072265625},{'x':628.0768432617188,'y':348.28363037109375},{'x':627.7057495117188,'y':347.9069519042969},{'x':627.3364868164062,'y':347.5285339355469},{'x':626.9689331054688,'y':347.14837646484375},{'x':626.603271484375,'y':346.7663879394531},{'x':626.2395629882812,'y':346.38262939453125},{'x':625.8778076171875,'y':345.9969482421875},{'x':625.518310546875,'y':345.6092224121094},{'x':625.1611328125,'y':345.21929931640625},{'x':624.8068237304688,'y':344.8268127441406},{'x':624.4552612304688,'y':344.4318542480469},{'x':624.1066284179688,'y':344.0342712402344},{'x':623.76123046875,'y':343.6339111328125},{'x':623.4193725585938,'y':343.2305603027344},{'x':623.0819091796875,'y':342.823486328125},{'x':622.7485961914062,'y':342.41302490234375},{'x':622.4198608398438,'y':341.9988708496094},{'x':622.0963745117188,'y':341.5805969238281},{'x':621.7791137695312,'y':341.1575622558594},{'x':621.4686889648438,'y':340.72955322265625},{'x':621.165771484375,'y':340.296142578125},{'x':620.8717041015625,'y':339.8567199707031},{'x':620.5880126953125,'y':339.4104919433594},{'x':620.31640625,'y':338.95684814453125},{'x':620.0592041015625,'y':338.494873046875},{'x':619.8197631835938,'y':338.02349853515625},{'x':619.60205078125,'y':337.5417175292969},{'x':619.4119873046875,'y':337.04840087890625},{'x':619.2584228515625,'y':336.5425720214844},{'x':619.15380859375,'y':336.0245361328125},{'x':619.1163940429688,'y':335.4975891113281},{'x':619.1705932617188,'y':334.9725341796875},{'x':619.3410034179688,'y':334.4734802246094},{'x':619.6347045898438,'y':334.03564453125},{'x':620.0236206054688,'y':333.6789245605469},{'x':620.4705810546875,'y':333.3973388671875},{'x':620.9505004882812,'y':333.17608642578125},{'x':621.4495239257812,'y':333.001708984375},{'x':621.9599609375,'y':332.864013671875},{'x':622.4772338867188,'y':332.7549133300781},{'x':622.9990234375,'y':332.6694030761719},{'x':623.5235595703125,'y':332.6031188964844},{'x':624.0499267578125,'y':332.55303955078125},{'x':624.5774536132812,'y':332.5166931152344},{'x':625.1056518554688,'y':332.4925231933594},{'x':625.6342163085938,'y':332.4786376953125},{'x':626.1629638671875,'y':332.4738464355469},{'x':626.6917114257812,'y':332.4774475097656},{'x':627.2203979492188,'y':332.4881896972656},{'x':627.7488403320312,'y':332.50579833984375},{'x':628.277099609375,'y':332.5293884277344},{'x':628.8050537109375,'y':332.55828857421875},{'x':629.3327026367188,'y':332.59210205078125},{'x':629.860107421875,'y':332.6304931640625},{'x':630.3871459960938,'y':332.6730651855469},{'x':630.913818359375,'y':332.71966552734375},{'x':631.440185546875,'y':332.7701416015625},{'x':631.9661865234375,'y':332.8240051269531},{'x':632.4918823242188,'y':332.8810729980469},{'x':633.0172119140625,'y':332.9410705566406},{'x':633.542236328125,'y':333.00384521484375},{'x':634.0669555664062,'y':333.0692443847656},{'x':634.5913696289062,'y':333.1371154785156},{'x':635.1154174804688,'y':333.20733642578125},{'x':635.6392211914062,'y':333.27978515625},{'x':636.1627197265625,'y':333.3542785644531},{'x':636.6859130859375,'y':333.43084716796875},{'x':637.208740234375,'y':333.5098571777344},{'x':637.7314453125,'y':333.589599609375},{'x':638.2538452148438,'y':333.6716003417969},{'x':638.7760009765625,'y':333.7548522949219},{'x':639.2979736328125,'y':333.8394775390625},{'x':639.8196411132812,'y':333.92578125},{'x':640.3411865234375,'y':334.0127868652344},{'x':640.8624267578125,'y':334.1016845703125},{'x':641.383544921875,'y':334.1914367675781},{'x':641.9043579101562,'y':334.2825622558594},{'x':642.4252319335938,'y':334.37371826171875},{'x':642.9459838867188,'y':334.46527099609375},{'x':643.466552734375,'y':334.5581970214844},{'x':643.987060546875,'y':334.651123046875},{'x':644.507568359375,'y':334.7441711425781},{'x':645.028076171875,'y':334.8374328613281},{'x':645.548583984375,'y':334.9306640625},{'x':646.0690307617188,'y':335.02392578125},{'x':646.5895385742188,'y':335.1171875},{'x':647.1099853515625,'y':335.21044921875},{'x':647.6304931640625,'y':335.3037109375},{'x':648.1509399414062,'y':335.3969421386719},{'x':648.6714477539062,'y':335.4902038574219},{'x':649.19189453125,'y':335.5834655761719},{'x':649.71240234375,'y':335.6767272949219},{'x':650.2328491210938,'y':335.7699890136719},{'x':650.7533569335938,'y':335.86322021484375},{'x':651.2738037109375,'y':335.95648193359375},{'x':651.7943115234375,'y':336.04974365234375}]]"
+        )
+
+        wantracktovideo = NODE_CLASS_MAPPINGS["WanTrackToVideo"]()
+        wantracktovideo_248 = wantracktovideo.encode(
+            tracks=get_value_at_index(primitivestringmultiline_247, 0),
+            width=832,
+            height=480,
+            length=81,
+            batch_size=1,
+            temperature=220,
+            topk=10,
+            positive=get_value_at_index(cliptextencode_6, 0),
+            negative=get_value_at_index(cliptextencode_7, 0),
+            vae=get_value_at_index(vaeloader_39, 0),
+            start_image=get_value_at_index(loadimage_240, 0),
+        )
+
+        modelsamplingsd3 = NODE_CLASS_MAPPINGS["ModelSamplingSD3"]()
+        ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
+        vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+        createvideo = NODE_CLASS_MAPPINGS["CreateVideo"]()
+        savevideo = NODE_CLASS_MAPPINGS["SaveVideo"]()
+
+        for q in range(1):
+            modelsamplingsd3_48 = modelsamplingsd3.patch(
+                shift=5, model=get_value_at_index(loraloadermodelonly_154, 0)
+            )
+
+            ksampler_3 = ksampler.sample(
+                seed=random.randint(1, 2**64),
+                steps=4,
+                cfg=1,
+                sampler_name="uni_pc",
+                scheduler="simple",
+                denoise=1,
+                model=get_value_at_index(modelsamplingsd3_48, 0),
+                positive=get_value_at_index(wantracktovideo_248, 0),
+                negative=get_value_at_index(wantracktovideo_248, 1),
+                latent_image=get_value_at_index(wantracktovideo_248, 2),
+            )
+
+            vaedecode_8 = vaedecode.decode(
+                samples=get_value_at_index(ksampler_3, 0),
+                vae=get_value_at_index(vaeloader_39, 0),
+            )
+
+            createvideo_68 = createvideo.create_video(
+                fps=16, images=get_value_at_index(vaedecode_8, 0)
+            )
+
+            savevideo_69 = savevideo.save_video(
+                filename_prefix="video/ComfyUI",
+                format="auto",
+                codec="auto",
+                video=get_value_at_index(createvideo_68, 0),
+            )
+
+
+if __name__ == "__main__":
+    main()