t# This is a combination of 2 commits.

d

.ci/windows_amd_base_files/README_VERY_IMPORTANT.txt

@@ -1,27 +0,0 @@
-As of the time of writing this you need this preview driver for best results:
-https://www.amd.com/en/resources/support-articles/release-notes/RN-AMDGPU-WINDOWS-PYTORCH-PREVIEW.html
-
-HOW TO RUN:
-
-If you have a AMD gpu:
-
-run_amd_gpu.bat
-
-If you have memory issues you can try disabling the smart memory management by running comfyui with:
-
-run_amd_gpu_disable_smart_memory.bat
-
-IF YOU GET A RED ERROR IN THE UI MAKE SURE YOU HAVE A MODEL/CHECKPOINT IN: ComfyUI\models\checkpoints
-
-You can download the stable diffusion XL one from: https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors
-
-
-RECOMMENDED WAY TO UPDATE:
-To update the ComfyUI code: update\update_comfyui.bat
-
-
-TO SHARE MODELS BETWEEN COMFYUI AND ANOTHER UI:
-In the ComfyUI directory you will find a file: extra_model_paths.yaml.example
-Rename this file to: extra_model_paths.yaml and edit it with your favorite text editor.
-
-

.ci/windows_amd_base_files/run_amd_gpu.bat

@@ -1,2 +0,0 @@
-.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build
-pause

.ci/windows_amd_base_files/run_amd_gpu_disable_smart_memory.bat

@@ -1,2 +0,0 @@
-.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --disable-smart-memory
-pause

.github/workflows/release-stable-all.yml

@@ -1,61 +0,0 @@
-name: "Release Stable All Portable Versions"
-
-on:
-  workflow_dispatch:
-    inputs:
-      git_tag:
-        description: 'Git tag'
-        required: true
-        type: string
-
-jobs:
-  release_nvidia_default:
-    permissions:
-      contents: "write"
-      packages: "write"
-      pull-requests: "read"
-    name: "Release NVIDIA Default (cu129)"
-    uses: ./.github/workflows/stable-release.yml
-    with:
-      git_tag: ${{ inputs.git_tag }}
-      cache_tag: "cu129"
-      python_minor: "13"
-      python_patch: "6"
-      rel_name: "nvidia"
-      rel_extra_name: ""
-      test_release: true
-    secrets: inherit
-
-  release_nvidia_cu128:
-    permissions:
-      contents: "write"
-      packages: "write"
-      pull-requests: "read"
-    name: "Release NVIDIA cu128"
-    uses: ./.github/workflows/stable-release.yml
-    with:
-      git_tag: ${{ inputs.git_tag }}
-      cache_tag: "cu128"
-      python_minor: "12"
-      python_patch: "10"
-      rel_name: "nvidia"
-      rel_extra_name: "_cu128"
-      test_release: true
-    secrets: inherit
-
-  release_amd_rocm:
-    permissions:
-      contents: "write"
-      packages: "write"
-      pull-requests: "read"
-    name: "Release AMD ROCm 6.4.4"
-    uses: ./.github/workflows/stable-release.yml
-    with:
-      git_tag: ${{ inputs.git_tag }}
-      cache_tag: "rocm644"
-      python_minor: "12"
-      python_patch: "10"
-      rel_name: "amd"
-      rel_extra_name: ""
-      test_release: false
-    secrets: inherit

.github/workflows/ruff.yml

@@ -21,28 +21,3 @@ jobs:
 
     - name: Run Ruff
       run: ruff check .
-
-  pylint:
-    name: Run Pylint
-    runs-on: ubuntu-latest
-
-    steps:
-    - name: Checkout repository
-      uses: actions/checkout@v4
-
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.12'
-
-    - name: Install requirements
-      run: |
-        python -m pip install --upgrade pip
-        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-        pip install -r requirements.txt
-
-    - name: Install Pylint
-      run: pip install pylint
-
-    - name: Run Pylint
-      run: pylint comfy_api_nodes

.github/workflows/stable-release.yml

@@ -2,53 +2,17 @@
 name: "Release Stable Version"
 
 on:
-  workflow_call:
-    inputs:
-      git_tag:
-        description: 'Git tag'
-        required: true
-        type: string
-      cache_tag:
-        description: 'Cached dependencies tag'
-        required: true
-        type: string
-        default: "cu129"
-      python_minor:
-        description: 'Python minor version'
-        required: true
-        type: string
-        default: "13"
-      python_patch:
-        description: 'Python patch version'
-        required: true
-        type: string
-        default: "6"
-      rel_name:
-        description: 'Release name'
-        required: true
-        type: string
-        default: "nvidia"
-      rel_extra_name:
-        description: 'Release extra name'
-        required: false
-        type: string
-        default: ""
-      test_release:
-        description: 'Test Release'
-        required: true
-        type: boolean
-        default: true
   workflow_dispatch:
     inputs:
       git_tag:
         description: 'Git tag'
         required: true
         type: string
-      cache_tag:
-        description: 'Cached dependencies tag'
+      cu:
+        description: 'CUDA version'
         required: true
         type: string
-        default: "cu129"
+        default: "129"
       python_minor:
         description: 'Python minor version'
         required: true
@@ -59,21 +23,7 @@ on:
         required: true
         type: string
         default: "6"
-      rel_name:
-        description: 'Release name'
-        required: true
-        type: string
-        default: "nvidia"
-      rel_extra_name:
-        description: 'Release extra name'
-        required: false
-        type: string
-        default: ""
-      test_release:
-        description: 'Test Release'
-        required: true
-        type: boolean
-        default: true
+
 
 jobs:
   package_comfy_windows:
@@ -92,15 +42,15 @@ jobs:
         id: cache
         with:
           path: |
-            ${{ inputs.cache_tag }}_python_deps.tar
+            cu${{ inputs.cu }}_python_deps.tar
             update_comfyui_and_python_dependencies.bat
-          key: ${{ runner.os }}-build-${{ inputs.cache_tag }}-${{ inputs.python_minor }}
+          key: ${{ runner.os }}-build-cu${{ inputs.cu }}-${{ inputs.python_minor }}
       - shell: bash
         run: |
-          mv ${{ inputs.cache_tag }}_python_deps.tar ../
+          mv cu${{ inputs.cu }}_python_deps.tar ../
           mv update_comfyui_and_python_dependencies.bat ../
           cd ..
-          tar xf ${{ inputs.cache_tag }}_python_deps.tar
+          tar xf cu${{ inputs.cu }}_python_deps.tar
           pwd
           ls
 
@@ -115,19 +65,12 @@ jobs:
           echo 'import site' >> ./python3${{ inputs.python_minor }}._pth
           curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
           ./python.exe get-pip.py
-          ./python.exe -s -m pip install ../${{ inputs.cache_tag }}_python_deps/*
-
-          grep comfyui ../ComfyUI/requirements.txt > ./requirements_comfyui.txt
-          ./python.exe -s -m pip install -r requirements_comfyui.txt
-          rm requirements_comfyui.txt
-
+          ./python.exe -s -m pip install ../cu${{ inputs.cu }}_python_deps/*
           sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
 
-          if test -f ./Lib/site-packages/torch/lib/dnnl.lib; then
-            rm ./Lib/site-packages/torch/lib/dnnl.lib #I don't think this is actually used and I need the space
-            rm ./Lib/site-packages/torch/lib/libprotoc.lib
-            rm ./Lib/site-packages/torch/lib/libprotobuf.lib
-          fi
+          rm ./Lib/site-packages/torch/lib/dnnl.lib #I don't think this is actually used and I need the space
+          rm ./Lib/site-packages/torch/lib/libprotoc.lib
+          rm ./Lib/site-packages/torch/lib/libprotobuf.lib
 
           cd ..
 
@@ -142,18 +85,14 @@ jobs:
 
           mkdir update
           cp -r ComfyUI/.ci/update_windows/* ./update/
-          cp -r ComfyUI/.ci/windows_${{ inputs.rel_name }}_base_files/* ./
+          cp -r ComfyUI/.ci/windows_base_files/* ./
           cp ../update_comfyui_and_python_dependencies.bat ./update/
 
           cd ..
 
           "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=768m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
-          mv ComfyUI_windows_portable.7z ComfyUI/ComfyUI_windows_portable_${{ inputs.rel_name }}${{ inputs.rel_extra_name }}.7z
+          mv ComfyUI_windows_portable.7z ComfyUI/ComfyUI_windows_portable_nvidia.7z
 
-      - shell: bash
-        if: ${{ inputs.test_release }}
-        run: |
-          cd ..
           cd ComfyUI_windows_portable
           python_embeded/python.exe -s ComfyUI/main.py --quick-test-for-ci --cpu
 
@@ -162,9 +101,10 @@ jobs:
           ls
 
       - name: Upload binaries to release
-        uses: softprops/action-gh-release@v2
+        uses: svenstaro/upload-release-action@v2
         with:
-          files: ComfyUI_windows_portable_${{ inputs.rel_name }}${{ inputs.rel_extra_name }}.7z
-          tag_name: ${{ inputs.git_tag }}
+          repo_token: ${{ secrets.GITHUB_TOKEN }}
+          file: ComfyUI_windows_portable_nvidia.7z
+          tag: ${{ inputs.git_tag }}
+          overwrite: true
           draft: true
-          overwrite_files: true

.github/workflows/test-unit.yml

@@ -10,7 +10,7 @@ jobs:
   test:
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-2022, macos-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
     runs-on: ${{ matrix.os }}
     continue-on-error: true
     steps:

.github/workflows/windows_release_dependencies.yml

@@ -56,8 +56,7 @@ jobs:
             ..\python_embeded\python.exe -s -m pip install --upgrade torch torchvision torchaudio ${{ inputs.xformers }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r ../ComfyUI/requirements.txt pygit2
             pause" > update_comfyui_and_python_dependencies.bat
 
-            grep -v comfyui requirements.txt > requirements_nocomfyui.txt
-            python -m pip wheel --no-cache-dir torch torchvision torchaudio ${{ inputs.xformers }} ${{ inputs.extra_dependencies }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r requirements_nocomfyui.txt pygit2 -w ./temp_wheel_dir
+            python -m pip wheel --no-cache-dir torch torchvision torchaudio ${{ inputs.xformers }} ${{ inputs.extra_dependencies }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r requirements.txt pygit2 -w ./temp_wheel_dir
             python -m pip install --no-cache-dir ./temp_wheel_dir/*
             echo installed basic
             ls -lah temp_wheel_dir

.github/workflows/windows_release_dependencies_manual.yml

@@ -1,64 +0,0 @@
-name: "Windows Release dependencies Manual"
-
-on:
-  workflow_dispatch:
-    inputs:
-      torch_dependencies:
-        description: 'torch dependencies'
-        required: false
-        type: string
-        default: "torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu128"
-      cache_tag:
-        description: 'Cached dependencies tag'
-        required: true
-        type: string
-        default: "cu128"
-
-      python_minor:
-        description: 'python minor version'
-        required: true
-        type: string
-        default: "12"
-
-      python_patch:
-        description: 'python patch version'
-        required: true
-        type: string
-        default: "10"
-
-jobs:
-  build_dependencies:
-    runs-on: windows-latest
-    steps:
-        - uses: actions/checkout@v4
-        - uses: actions/setup-python@v5
-          with:
-            python-version: 3.${{ inputs.python_minor }}.${{ inputs.python_patch }}
-
-        - shell: bash
-          run: |
-            echo "@echo off
-            call update_comfyui.bat nopause
-            echo -
-            echo This will try to update pytorch and all python dependencies.
-            echo -
-            echo If you just want to update normally, close this and run update_comfyui.bat instead.
-            echo -
-            pause
-            ..\python_embeded\python.exe -s -m pip install --upgrade ${{ inputs.torch_dependencies }} -r ../ComfyUI/requirements.txt pygit2
-            pause" > update_comfyui_and_python_dependencies.bat
-
-            grep -v comfyui requirements.txt > requirements_nocomfyui.txt
-            python -m pip wheel --no-cache-dir ${{ inputs.torch_dependencies }} -r requirements_nocomfyui.txt pygit2 -w ./temp_wheel_dir
-            python -m pip install --no-cache-dir ./temp_wheel_dir/*
-            echo installed basic
-            ls -lah temp_wheel_dir
-            mv temp_wheel_dir ${{ inputs.cache_tag }}_python_deps
-            tar cf ${{ inputs.cache_tag }}_python_deps.tar ${{ inputs.cache_tag }}_python_deps
-
-        - uses: actions/cache/save@v4
-          with:
-            path: |
-              ${{ inputs.cache_tag }}_python_deps.tar
-              update_comfyui_and_python_dependencies.bat
-            key: ${{ runner.os }}-build-${{ inputs.cache_tag }}-${{ inputs.python_minor }}

.github/workflows/windows_release_nightly_pytorch.yml

@@ -68,7 +68,7 @@ jobs:
 
             mkdir update
             cp -r ComfyUI/.ci/update_windows/* ./update/
-            cp -r ComfyUI/.ci/windows_nvidia_base_files/* ./
+            cp -r ComfyUI/.ci/windows_base_files/* ./
             cp -r ComfyUI/.ci/windows_nightly_base_files/* ./
 
             echo "call update_comfyui.bat nopause

.github/workflows/windows_release_package.yml

@@ -81,7 +81,7 @@ jobs:
 
             mkdir update
             cp -r ComfyUI/.ci/update_windows/* ./update/
-            cp -r ComfyUI/.ci/windows_nvidia_base_files/* ./
+            cp -r ComfyUI/.ci/windows_base_files/* ./
             cp ../update_comfyui_and_python_dependencies.bat ./update/
 
             cd ..

CODEOWNERS

@@ -1,3 +1,25 @@
 # Admins
 * @comfyanonymous
-* @kosinkadink
+
+# Note: Github teams syntax cannot be used here as the repo is not owned by Comfy-Org.
+# Inlined the team members for now.
+
+# Maintainers
+*.md @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/tests/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/tests-unit/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/notebooks/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/script_examples/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/.github/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/requirements.txt @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/pyproject.toml @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+
+# Python web server
+/api_server/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+/app/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+/utils/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+
+# Node developers
+/comfy_extras/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill
+/comfy/comfy_types/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill
+/comfy_api_nodes/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill

README.md

@@ -176,12 +176,6 @@ Simply download, extract with [7-Zip](https://7-zip.org) and run. Make sure you
 
 If you have trouble extracting it, right click the file -> properties -> unblock
 
-#### Alternative Downloads:
-
-[Experimental portable for AMD GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_amd.7z)
-
-[Portable with pytorch cuda 12.8 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu128.7z) (Supports Nvidia 10 series and older GPUs).
-
 #### How do I share models between another UI and ComfyUI?
 
 See the [Config file](extra_model_paths.yaml.example) to set the search paths for models. In the standalone windows build you can find this file in the ComfyUI directory. Rename this file to extra_model_paths.yaml and edit it with your favorite text editor.
@@ -206,32 +200,14 @@ Put your SD checkpoints (the huge ckpt/safetensors files) in: models/checkpoints
 Put your VAE in: models/vae
 
 
-### AMD GPUs (Linux)
-
+### AMD GPUs (Linux only)
 AMD users can install rocm and pytorch with pip if you don't have it already installed, this is the command to install the stable version:
 
 ```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4```
 
-This is the command to install the nightly with ROCm 7.0 which might have some performance improvements:
+This is the command to install the nightly with ROCm 6.4 which might have some performance improvements:
 
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.0```
-
-
-### AMD GPUs (Experimental: Windows and Linux), RDNA 3, 3.5 and 4 only.
-
-These have less hardware support than the builds above but they work on windows. You also need to install the pytorch version specific to your hardware.
-
-RDNA 3 (RX 7000 series):
-
-```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/```
-
-RDNA 3.5 (Strix halo/Ryzen AI Max+ 365):
-
-```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx1151/```
-
-RDNA 4 (RX 9000 series):
-
-```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx120X-all/```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.4```
 
 ### Intel GPUs (Windows and Linux)
 
@@ -257,7 +233,7 @@ Nvidia users should install stable pytorch using this command:
 
 This is the command to install pytorch nightly instead which might have performance improvements.
 
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu129```
 
 #### Troubleshooting
 
@@ -288,6 +264,12 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve
 
 > **Note**: Remember to add your models, VAE, LoRAs etc. to the corresponding Comfy folders, as discussed in [ComfyUI manual installation](#manual-install-windows-linux).
 
+#### DirectML (AMD Cards on Windows)
+
+This is very badly supported and is not recommended. There are some unofficial builds of pytorch ROCm on windows that exist that will give you a much better experience than this. This readme will be updated once official pytorch ROCm builds for windows come out.
+
+```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml```
+
 #### Ascend NPUs
 
 For models compatible with Ascend Extension for PyTorch (torch_npu). To get started, ensure your environment meets the prerequisites outlined on the [installation](https://ascend.github.io/docs/sources/ascend/quick_install.html) page. Here's a step-by-step guide tailored to your platform and installation method:

app/frontend_management.py

@@ -42,7 +42,6 @@ def get_installed_frontend_version():
     frontend_version_str = version("comfyui-frontend-package")
     return frontend_version_str
 
-
 def get_required_frontend_version():
     """Get the required frontend version from requirements.txt."""
     try:
@@ -64,7 +63,6 @@ def get_required_frontend_version():
         logging.error(f"Error reading requirements.txt: {e}")
         return None
 
-
 def check_frontend_version():
     """Check if the frontend version is up to date."""
 
@@ -205,37 +203,6 @@ class FrontendManager:
         """Get the required frontend package version."""
         return get_required_frontend_version()
 
-    @classmethod
-    def get_installed_templates_version(cls) -> str:
-        """Get the currently installed workflow templates package version."""
-        try:
-            templates_version_str = version("comfyui-workflow-templates")
-            return templates_version_str
-        except Exception:
-            return None
-
-    @classmethod
-    def get_required_templates_version(cls) -> str:
-        """Get the required workflow templates version from requirements.txt."""
-        try:
-            with open(requirements_path, "r", encoding="utf-8") as f:
-                for line in f:
-                    line = line.strip()
-                    if line.startswith("comfyui-workflow-templates=="):
-                        version_str = line.split("==")[-1]
-                        if not is_valid_version(version_str):
-                            logging.error(f"Invalid templates version format in requirements.txt: {version_str}")
-                            return None
-                        return version_str
-                logging.error("comfyui-workflow-templates not found in requirements.txt")
-                return None
-        except FileNotFoundError:
-            logging.error("requirements.txt not found. Cannot determine required templates version.")
-            return None
-        except Exception as e:
-            logging.error(f"Error reading requirements.txt: {e}")
-            return None
-
     @classmethod
     def default_frontend_path(cls) -> str:
         try:

comfy/ldm/ace/vae/music_dcae_pipeline.py

@@ -23,6 +23,8 @@ class MusicDCAE(torch.nn.Module):
         else:
             self.source_sample_rate = source_sample_rate
 
+        # self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
+
         self.transform = transforms.Compose([
             transforms.Normalize(0.5, 0.5),
         ])
@@ -35,6 +37,10 @@ class MusicDCAE(torch.nn.Module):
         self.scale_factor = 0.1786
         self.shift_factor = -1.9091
 
+    def load_audio(self, audio_path):
+        audio, sr = torchaudio.load(audio_path)
+        return audio, sr
+
     def forward_mel(self, audios):
         mels = []
         for i in range(len(audios)):
@@ -67,8 +73,10 @@ class MusicDCAE(torch.nn.Module):
             latent = self.dcae.encoder(mel.unsqueeze(0))
             latents.append(latent)
         latents = torch.cat(latents, dim=0)
+        # latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
         latents = (latents - self.shift_factor) * self.scale_factor
         return latents
+        # return latents, latent_lengths
 
     @torch.no_grad()
     def decode(self, latents, audio_lengths=None, sr=None):
@@ -83,7 +91,9 @@ class MusicDCAE(torch.nn.Module):
             wav = self.vocoder.decode(mels[0]).squeeze(1)
 
             if sr is not None:
+                # resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
                 wav = torchaudio.functional.resample(wav, 44100, sr)
+                # wav = resampler(wav)
             else:
                 sr = 44100
             pred_wavs.append(wav)
@@ -91,6 +101,7 @@ class MusicDCAE(torch.nn.Module):
         if audio_lengths is not None:
             pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
         return torch.stack(pred_wavs)
+        # return sr, pred_wavs
 
     def forward(self, audios, audio_lengths=None, sr=None):
         latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)

comfy/ldm/flux/math.py

@@ -37,10 +37,7 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
 
 def apply_rope1(x: Tensor, freqs_cis: Tensor):
     x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
-
-    x_out = freqs_cis[..., 0] * x_[..., 0]
-    x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
-
+    x_out = freqs_cis[..., 0] * x_[..., 0] + freqs_cis[..., 1] * x_[..., 1]
     return x_out.reshape(*x.shape).type_as(x)
 
 def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):

comfy/ldm/hunyuan_video/vae_refiner.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, VideoConv3d, Normalize
+from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, VideoConv3d
 import comfy.ops
 import comfy.ldm.models.autoencoder
 ops = comfy.ops.disable_weight_init
@@ -17,12 +17,11 @@ class RMS_norm(nn.Module):
         return F.normalize(x, dim=1) * self.scale * self.gamma
 
 class DnSmpl(nn.Module):
-    def __init__(self, ic, oc, tds=True, refiner_vae=True, op=VideoConv3d):
+    def __init__(self, ic, oc, tds=True):
         super().__init__()
         fct = 2 * 2 * 2 if tds else 1 * 2 * 2
         assert oc % fct == 0
-        self.conv = op(ic, oc // fct, kernel_size=3, stride=1, padding=1)
-        self.refiner_vae = refiner_vae
+        self.conv = VideoConv3d(ic, oc // fct, kernel_size=3)
 
         self.tds = tds
         self.gs = fct * ic // oc
@@ -31,7 +30,7 @@ class DnSmpl(nn.Module):
         r1 = 2 if self.tds else 1
         h = self.conv(x)
 
-        if self.tds and self.refiner_vae:
+        if self.tds:
             hf = h[:, :, :1, :, :]
             b, c, f, ht, wd = hf.shape
             hf = hf.reshape(b, c, f, ht // 2, 2, wd // 2, 2)
@@ -67,7 +66,6 @@ class DnSmpl(nn.Module):
             sc = torch.cat([xf, xn], dim=2)
         else:
             b, c, frms, ht, wd = h.shape
-
             nf = frms // r1
             h = h.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
             h = h.permute(0, 3, 5, 7, 1, 2, 4, 6)
@@ -85,11 +83,10 @@ class DnSmpl(nn.Module):
 
 
 class UpSmpl(nn.Module):
-    def __init__(self, ic, oc, tus=True, refiner_vae=True, op=VideoConv3d):
+    def __init__(self, ic, oc, tus=True):
         super().__init__()
         fct = 2 * 2 * 2 if tus else 1 * 2 * 2
-        self.conv = op(ic, oc * fct, kernel_size=3, stride=1, padding=1)
-        self.refiner_vae = refiner_vae
+        self.conv = VideoConv3d(ic, oc * fct, kernel_size=3)
 
         self.tus = tus
         self.rp = fct * oc // ic
@@ -98,7 +95,7 @@ class UpSmpl(nn.Module):
         r1 = 2 if self.tus else 1
         h = self.conv(x)
 
-        if self.tus and self.refiner_vae:
+        if self.tus:
             hf = h[:, :, :1, :, :]
             b, c, f, ht, wd = hf.shape
             nc = c // (2 * 2)
@@ -151,56 +148,43 @@ class UpSmpl(nn.Module):
 
 class Encoder(nn.Module):
     def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, ffactor_temporal, downsample_match_channel=True, refiner_vae=True, **_):
+                 ffactor_spatial, ffactor_temporal, downsample_match_channel=True, **_):
         super().__init__()
         self.z_channels = z_channels
         self.block_out_channels = block_out_channels
         self.num_res_blocks = num_res_blocks
-        self.ffactor_temporal = ffactor_temporal
-
-        self.refiner_vae = refiner_vae
-        if self.refiner_vae:
-            conv_op = VideoConv3d
-            norm_op = RMS_norm
-        else:
-            conv_op = ops.Conv3d
-            norm_op = Normalize
-
-        self.conv_in = conv_op(in_channels, block_out_channels[0], 3, 1, 1)
+        self.conv_in = VideoConv3d(in_channels, block_out_channels[0], 3, 1, 1)
 
         self.down = nn.ModuleList()
         ch = block_out_channels[0]
         depth = (ffactor_spatial >> 1).bit_length()
-        depth_temporal = ((ffactor_spatial // self.ffactor_temporal) >> 1).bit_length()
+        depth_temporal = ((ffactor_spatial // ffactor_temporal) >> 1).bit_length()
 
         for i, tgt in enumerate(block_out_channels):
             stage = nn.Module()
             stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
                                                      out_channels=tgt,
                                                      temb_channels=0,
-                                                     conv_op=conv_op, norm_op=norm_op)
+                                                     conv_op=VideoConv3d, norm_op=RMS_norm)
                                         for j in range(num_res_blocks)])
             ch = tgt
             if i < depth:
                 nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and downsample_match_channel else ch
-                stage.downsample = DnSmpl(ch, nxt, tds=i >= depth_temporal, refiner_vae=self.refiner_vae, op=conv_op)
+                stage.downsample = DnSmpl(ch, nxt, tds=i >= depth_temporal)
                 ch = nxt
             self.down.append(stage)
 
         self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
+        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=RMS_norm)
+        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
 
-        self.norm_out = norm_op(ch)
-        self.conv_out = conv_op(ch, z_channels << 1, 3, 1, 1)
+        self.norm_out = RMS_norm(ch)
+        self.conv_out = VideoConv3d(ch, z_channels << 1, 3, 1, 1)
 
         self.regul = comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer()
 
     def forward(self, x):
-        if not self.refiner_vae and x.shape[2] == 1:
-            x = x.expand(-1, -1, self.ffactor_temporal, -1, -1)
-
         x = self.conv_in(x)
 
         for stage in self.down:
@@ -216,42 +200,31 @@ class Encoder(nn.Module):
         skip = x.view(b, c // grp, grp, t, h, w).mean(2)
 
         out = self.conv_out(F.silu(self.norm_out(x))) + skip
+        out = self.regul(out)[0]
 
-        if self.refiner_vae:
-            out = self.regul(out)[0]
-
-            out = torch.cat((out[:, :, :1], out), dim=2)
-            out = out.permute(0, 2, 1, 3, 4)
-            b, f_times_2, c, h, w = out.shape
-            out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
-            out = out.permute(0, 2, 1, 3, 4).contiguous()
-
+        out = torch.cat((out[:, :, :1], out), dim=2)
+        out = out.permute(0, 2, 1, 3, 4)
+        b, f_times_2, c, h, w = out.shape
+        out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
+        out = out.permute(0, 2, 1, 3, 4).contiguous()
         return out
 
 class Decoder(nn.Module):
     def __init__(self, z_channels, out_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, ffactor_temporal, upsample_match_channel=True, refiner_vae=True, **_):
+                 ffactor_spatial, ffactor_temporal, upsample_match_channel=True, **_):
         super().__init__()
         block_out_channels = block_out_channels[::-1]
         self.z_channels = z_channels
         self.block_out_channels = block_out_channels
         self.num_res_blocks = num_res_blocks
 
-        self.refiner_vae = refiner_vae
-        if self.refiner_vae:
-            conv_op = VideoConv3d
-            norm_op = RMS_norm
-        else:
-            conv_op = ops.Conv3d
-            norm_op = Normalize
-
         ch = block_out_channels[0]
-        self.conv_in = conv_op(z_channels, ch, kernel_size=3, stride=1, padding=1)
+        self.conv_in = VideoConv3d(z_channels, ch, 3)
 
         self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
+        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=RMS_norm)
+        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
 
         self.up = nn.ModuleList()
         depth = (ffactor_spatial >> 1).bit_length()
@@ -262,26 +235,25 @@ class Decoder(nn.Module):
             stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
                                                      out_channels=tgt,
                                                      temb_channels=0,
-                                                     conv_op=conv_op, norm_op=norm_op)
+                                                     conv_op=VideoConv3d, norm_op=RMS_norm)
                                         for j in range(num_res_blocks + 1)])
             ch = tgt
             if i < depth:
                 nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and upsample_match_channel else ch
-                stage.upsample = UpSmpl(ch, nxt, tus=i < depth_temporal, refiner_vae=self.refiner_vae, op=conv_op)
+                stage.upsample = UpSmpl(ch, nxt, tus=i < depth_temporal)
                 ch = nxt
             self.up.append(stage)
 
-        self.norm_out = norm_op(ch)
-        self.conv_out = conv_op(ch, out_channels, 3, stride=1, padding=1)
+        self.norm_out = RMS_norm(ch)
+        self.conv_out = VideoConv3d(ch, out_channels, 3)
 
     def forward(self, z):
-        if self.refiner_vae:
-            z = z.permute(0, 2, 1, 3, 4)
-            b, f, c, h, w = z.shape
-            z = z.reshape(b, f, 2, c // 2, h, w)
-            z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
-            z = z.permute(0, 2, 1, 3, 4)
-            z = z[:, :, 1:]
+        z = z.permute(0, 2, 1, 3, 4)
+        b, f, c, h, w = z.shape
+        z = z.reshape(b, f, 2, c // 2, h, w)
+        z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
+        z = z.permute(0, 2, 1, 3, 4)
+        z = z[:, :, 1:]
 
         x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
         x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
@@ -292,10 +264,4 @@ class Decoder(nn.Module):
             if hasattr(stage, 'upsample'):
                 x = stage.upsample(x)
 
-        out = self.conv_out(F.silu(self.norm_out(x)))
-
-        if not self.refiner_vae:
-            if z.shape[-3] == 1:
-                out = out[:, :, -1:]
-
-        return out
+        return self.conv_out(F.silu(self.norm_out(x)))

comfy/ldm/wan/model.py

@@ -237,7 +237,6 @@ class WanAttentionBlock(nn.Module):
             freqs, transformer_options=transformer_options)
 
         x = torch.addcmul(x, y, repeat_e(e[2], x))
-        del y
 
         # cross-attention & ffn
         x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
@@ -903,7 +902,7 @@ class MotionEncoder_tc(nn.Module):
     def __init__(self,
                  in_dim: int,
                  hidden_dim: int,
-                 num_heads: int,
+                 num_heads=int,
                  need_global=True,
                  dtype=None,
                  device=None,
@@ -1356,7 +1355,7 @@ class WanT2VCrossAttentionGather(WanSelfAttention):
 
         x = optimized_attention(q, k, v, heads=self.num_heads, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
 
-        x = x.transpose(1, 2).reshape(b, -1, n * d)
+        x = x.transpose(1, 2).view(b, -1, n, d).flatten(2)
         x = self.o(x)
         return x
 

comfy/ldm/wan/vae.py

@@ -468,46 +468,55 @@ class WanVAE(nn.Module):
                                  attn_scales, self.temperal_upsample, dropout)
 
     def encode(self, x):
-        conv_idx = [0]
-        feat_map = [None] * count_conv3d(self.decoder)
+        self.clear_cache()
         ## cache
         t = x.shape[2]
         iter_ = 1 + (t - 1) // 4
         ## 对encode输入的x，按时间拆分为1、4、4、4....
         for i in range(iter_):
-            conv_idx = [0]
+            self._enc_conv_idx = [0]
             if i == 0:
                 out = self.encoder(
                     x[:, :, :1, :, :],
-                    feat_cache=feat_map,
-                    feat_idx=conv_idx)
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
             else:
                 out_ = self.encoder(
                     x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
-                    feat_cache=feat_map,
-                    feat_idx=conv_idx)
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
                 out = torch.cat([out, out_], 2)
         mu, log_var = self.conv1(out).chunk(2, dim=1)
+        self.clear_cache()
         return mu
 
     def decode(self, z):
-        conv_idx = [0]
-        feat_map = [None] * count_conv3d(self.decoder)
+        self.clear_cache()
         # z: [b,c,t,h,w]
 
         iter_ = z.shape[2]
         x = self.conv2(z)
         for i in range(iter_):
-            conv_idx = [0]
+            self._conv_idx = [0]
             if i == 0:
                 out = self.decoder(
                     x[:, :, i:i + 1, :, :],
-                    feat_cache=feat_map,
-                    feat_idx=conv_idx)
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
             else:
                 out_ = self.decoder(
                     x[:, :, i:i + 1, :, :],
-                    feat_cache=feat_map,
-                    feat_idx=conv_idx)
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
                 out = torch.cat([out, out_], 2)
+        self.clear_cache()
         return out
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        #cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num

comfy/model_detection.py

@@ -365,8 +365,8 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["patch_size"] = 2
         dit_config["in_channels"] = 16
         dit_config["dim"] = 2304
-        dit_config["cap_feat_dim"] = state_dict['{}cap_embedder.1.weight'.format(key_prefix)].shape[1]
-        dit_config["n_layers"] = count_blocks(state_dict_keys, '{}layers.'.format(key_prefix) + '{}.')
+        dit_config["cap_feat_dim"] = 2304
+        dit_config["n_layers"] = 26
         dit_config["n_heads"] = 24
         dit_config["n_kv_heads"] = 8
         dit_config["qk_norm"] = True

comfy/model_management.py

@@ -645,9 +645,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
             if loaded_model.model.is_clone(current_loaded_models[i].model):
                 to_unload = [i] + to_unload
         for i in to_unload:
-            model_to_unload = current_loaded_models.pop(i)
-            model_to_unload.model.detach(unpatch_all=False)
-            model_to_unload.model_finalizer.detach()
+            current_loaded_models.pop(i).model.detach(unpatch_all=False)
 
     total_memory_required = {}
     for loaded_model in models_to_load:

comfy/model_patcher.py

@@ -123,30 +123,16 @@ def move_weight_functions(m, device):
     return memory
 
 class LowVramPatch:
-    def __init__(self, key, patches, convert_func=None, set_func=None):
+    def __init__(self, key, patches):
         self.key = key
         self.patches = patches
-        self.convert_func = convert_func
-        self.set_func = set_func
-
     def __call__(self, weight):
         intermediate_dtype = weight.dtype
-        if self.convert_func is not None:
-            weight = self.convert_func(weight.to(dtype=torch.float32, copy=True), inplace=True)
-
         if intermediate_dtype not in [torch.float32, torch.float16, torch.bfloat16]: #intermediate_dtype has to be one that is supported in math ops
             intermediate_dtype = torch.float32
-            out = comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype)
-            if self.set_func is None:
-                return comfy.float.stochastic_rounding(out, weight.dtype, seed=string_to_seed(self.key))
-            else:
-                return self.set_func(out, seed=string_to_seed(self.key), return_weight=True)
+            return comfy.float.stochastic_rounding(comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype), weight.dtype, seed=string_to_seed(self.key))
 
-        out = comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
-        if self.set_func is not None:
-            return self.set_func(out, seed=string_to_seed(self.key), return_weight=True).to(dtype=intermediate_dtype)
-        else:
-            return out
+        return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
 
 def get_key_weight(model, key):
     set_func = None
@@ -671,15 +657,13 @@ class ModelPatcher:
                         if force_patch_weights:
                             self.patch_weight_to_device(weight_key)
                         else:
-                            _, set_func, convert_func = get_key_weight(self.model, weight_key)
-                            m.weight_function = [LowVramPatch(weight_key, self.patches, convert_func, set_func)]
+                            m.weight_function = [LowVramPatch(weight_key, self.patches)]
                             patch_counter += 1
                     if bias_key in self.patches:
                         if force_patch_weights:
                             self.patch_weight_to_device(bias_key)
                         else:
-                            _, set_func, convert_func = get_key_weight(self.model, bias_key)
-                            m.bias_function = [LowVramPatch(bias_key, self.patches, convert_func, set_func)]
+                            m.bias_function = [LowVramPatch(bias_key, self.patches)]
                             patch_counter += 1
 
                     cast_weight = True
@@ -841,12 +825,10 @@ class ModelPatcher:
                         module_mem += move_weight_functions(m, device_to)
                         if lowvram_possible:
                             if weight_key in self.patches:
-                                _, set_func, convert_func = get_key_weight(self.model, weight_key)
-                                m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
+                                m.weight_function.append(LowVramPatch(weight_key, self.patches))
                                 patch_counter += 1
                             if bias_key in self.patches:
-                                _, set_func, convert_func = get_key_weight(self.model, bias_key)
-                                m.bias_function.append(LowVramPatch(bias_key, self.patches, convert_func, set_func))
+                                m.bias_function.append(LowVramPatch(bias_key, self.patches))
                                 patch_counter += 1
                             cast_weight = True
 

comfy/model_sampling.py

@@ -21,23 +21,17 @@ def rescale_zero_terminal_snr_sigmas(sigmas):
     alphas_bar[-1] = 4.8973451890853435e-08
     return ((1 - alphas_bar) / alphas_bar) ** 0.5
 
-def reshape_sigma(sigma, noise_dim):
-    if sigma.nelement() == 1:
-        return sigma.view(())
-    else:
-        return sigma.view(sigma.shape[:1] + (1,) * (noise_dim - 1))
-
 class EPS:
     def calculate_input(self, sigma, noise):
-        sigma = reshape_sigma(sigma, noise.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         return noise / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
 
     def calculate_denoised(self, sigma, model_output, model_input):
-        sigma = reshape_sigma(sigma, model_output.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
         return model_input - model_output * sigma
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
-        sigma = reshape_sigma(sigma, noise.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         if max_denoise:
             noise = noise * torch.sqrt(1.0 + sigma ** 2.0)
         else:
@@ -51,12 +45,12 @@ class EPS:
 
 class V_PREDICTION(EPS):
     def calculate_denoised(self, sigma, model_output, model_input):
-        sigma = reshape_sigma(sigma, model_output.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
         return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) - model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
 
 class EDM(V_PREDICTION):
     def calculate_denoised(self, sigma, model_output, model_input):
-        sigma = reshape_sigma(sigma, model_output.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
         return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) + model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
 
 class CONST:
@@ -64,15 +58,15 @@ class CONST:
         return noise
 
     def calculate_denoised(self, sigma, model_output, model_input):
-        sigma = reshape_sigma(sigma, model_output.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
         return model_input - model_output * sigma
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
-        sigma = reshape_sigma(sigma, noise.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         return sigma * noise + (1.0 - sigma) * latent_image
 
     def inverse_noise_scaling(self, sigma, latent):
-        sigma = reshape_sigma(sigma, latent.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (latent.ndim - 1))
         return latent / (1.0 - sigma)
 
 class X0(EPS):
@@ -86,16 +80,16 @@ class IMG_TO_IMG(X0):
 class COSMOS_RFLOW:
     def calculate_input(self, sigma, noise):
         sigma = (sigma / (sigma + 1))
-        sigma = reshape_sigma(sigma, noise.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         return noise * (1.0 - sigma)
 
     def calculate_denoised(self, sigma, model_output, model_input):
         sigma = (sigma / (sigma + 1))
-        sigma = reshape_sigma(sigma, model_output.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
         return model_input * (1.0 - sigma) - model_output * sigma
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
-        sigma = reshape_sigma(sigma, noise.ndim)
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         noise = noise * sigma
         noise += latent_image
         return noise

comfy/ops.py

@@ -416,10 +416,8 @@ def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None
                 else:
                     return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype)
 
-            def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
+            def set_weight(self, weight, inplace_update=False, seed=None, **kwargs):
                 weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed)
-                if return_weight:
-                    return weight
                 if inplace_update:
                     self.weight.data.copy_(weight)
                 else:

comfy/samplers.py

@@ -360,7 +360,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
 def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
     if "sampler_cfg_function" in model_options:
         args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
-                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options, "input_cond": cond, "input_uncond": uncond}
+                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
         cfg_result = x - model_options["sampler_cfg_function"](args)
     else:
         cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
@@ -390,7 +390,7 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option
     for fn in model_options.get("sampler_pre_cfg_function", []):
         args = {"conds":conds, "conds_out": out, "cond_scale": cond_scale, "timestep": timestep,
                 "input": x, "sigma": timestep, "model": model, "model_options": model_options}
-        out = fn(args)
+        out  = fn(args)
 
     return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
 

comfy/sd.py

@@ -332,51 +332,35 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
+            elif "decoder.conv_in.weight" in sd and sd['decoder.conv_in.weight'].shape[1] == 64:
+                ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
+                self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
+                self.downscale_ratio = 32
+                self.upscale_ratio = 32
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
+                                                            encoder_config={'target': "comfy.ldm.hunyuan_video.vae.Encoder", 'params': ddconfig},
+                                                            decoder_config={'target': "comfy.ldm.hunyuan_video.vae.Decoder", 'params': ddconfig})
+
+                self.memory_used_encode = lambda shape, dtype: (700 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.memory_used_decode = lambda shape, dtype: (700 * shape[2] * shape[3] * 32 * 32) * model_management.dtype_size(dtype)
+
             elif "decoder.conv_in.weight" in sd:
-                if sd['decoder.conv_in.weight'].shape[1] == 64:
-                    ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
-                    self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                    self.downscale_ratio = 32
-                    self.upscale_ratio = 32
-                    self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-                    self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
-                                                                encoder_config={'target': "comfy.ldm.hunyuan_video.vae.Encoder", 'params': ddconfig},
-                                                                decoder_config={'target': "comfy.ldm.hunyuan_video.vae.Decoder", 'params': ddconfig})
+                #default SD1.x/SD2.x VAE parameters
+                ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
 
-                    self.memory_used_encode = lambda shape, dtype: (700 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
-                    self.memory_used_decode = lambda shape, dtype: (700 * shape[2] * shape[3] * 32 * 32) * model_management.dtype_size(dtype)
-                elif sd['decoder.conv_in.weight'].shape[1] == 32:
-                    ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True, "refiner_vae": False}
-                    self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                    self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-                    self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
-                    self.upscale_index_formula = (4, 16, 16)
-                    self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
-                    self.downscale_index_formula = (4, 16, 16)
-                    self.latent_dim = 3
-                    self.not_video = True
-                    self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
-                                                                encoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Encoder", 'params': ddconfig},
-                                                                decoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Decoder", 'params': ddconfig})
+                if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
+                    ddconfig['ch_mult'] = [1, 2, 4]
+                    self.downscale_ratio = 4
+                    self.upscale_ratio = 4
 
-                    self.memory_used_encode = lambda shape, dtype: (2800 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
-                    self.memory_used_decode = lambda shape, dtype: (2800 * shape[-3] * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
+                self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
+                if 'post_quant_conv.weight' in sd:
+                    self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
                 else:
-                    #default SD1.x/SD2.x VAE parameters
-                    ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
-
-                    if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
-                        ddconfig['ch_mult'] = [1, 2, 4]
-                        self.downscale_ratio = 4
-                        self.upscale_ratio = 4
-
-                    self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                    if 'post_quant_conv.weight' in sd:
-                        self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
-                    else:
-                        self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
-                                                                    encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
-                                                                    decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
+                    self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
+                                                                encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
+                                                                decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
             elif "decoder.layers.1.layers.0.beta" in sd:
                 self.first_stage_model = AudioOobleckVAE()
                 self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)
@@ -652,7 +636,6 @@ class VAE:
     def decode(self, samples_in, vae_options={}):
         self.throw_exception_if_invalid()
         pixel_samples = None
-        do_tile = False
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -668,13 +651,6 @@ class VAE:
                 pixel_samples[x:x+batch_number] = out
         except model_management.OOM_EXCEPTION:
             logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
-            #NOTE: We don't know what tensors were allocated to stack variables at the time of the
-            #exception and the exception itself refs them all until we get out of this except block.
-            #So we just set a flag for tiler fallback so that tensor gc can happen once the
-            #exception is fully off the books.
-            do_tile = True
-
-        if do_tile:
             dims = samples_in.ndim - 2
             if dims == 1 or self.extra_1d_channel is not None:
                 pixel_samples = self.decode_tiled_1d(samples_in)
@@ -721,7 +697,6 @@ class VAE:
         self.throw_exception_if_invalid()
         pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         pixel_samples = pixel_samples.movedim(-1, 1)
-        do_tile = False
         if self.latent_dim == 3 and pixel_samples.ndim < 5:
             if not self.not_video:
                 pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
@@ -743,13 +718,6 @@ class VAE:
 
         except model_management.OOM_EXCEPTION:
             logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
-            #NOTE: We don't know what tensors were allocated to stack variables at the time of the
-            #exception and the exception itself refs them all until we get out of this except block.
-            #So we just set a flag for tiler fallback so that tensor gc can happen once the
-            #exception is fully off the books.
-            do_tile = True
-
-        if do_tile:
             if self.latent_dim == 3:
                 tile = 256
                 overlap = tile // 4
@@ -890,7 +858,6 @@ class TEModel(Enum):
     QWEN25_3B = 10
     QWEN25_7B = 11
     BYT5_SMALL_GLYPH = 12
-    GEMMA_3_4B = 13
 
 def detect_te_model(sd):
     if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:
@@ -913,8 +880,6 @@ def detect_te_model(sd):
             return TEModel.BYT5_SMALL_GLYPH
         return TEModel.T5_BASE
     if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
-        if 'model.layers.0.self_attn.q_norm.weight' in sd:
-            return TEModel.GEMMA_3_4B
         return TEModel.GEMMA_2_2B
     if 'model.layers.0.self_attn.k_proj.bias' in sd:
         weight = sd['model.layers.0.self_attn.k_proj.bias']
@@ -1019,10 +984,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.lumina2.LuminaTokenizer
             tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
-        elif te_model == TEModel.GEMMA_3_4B:
-            clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data), model_type="gemma3_4b")
-            clip_target.tokenizer = comfy.text_encoders.lumina2.NTokenizer
-            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
         elif te_model == TEModel.LLAMA3_8:
             clip_target.clip = comfy.text_encoders.hidream.hidream_clip(**llama_detect(clip_data),
                                                                         clip_l=False, clip_g=False, t5=False, llama=True, dtype_t5=None, t5xxl_scaled_fp8=None)

comfy/text_encoders/hunyuan_image.py

@@ -63,13 +63,7 @@ class HunyuanImageTEModel(QwenImageTEModel):
             self.byt5_small = None
 
     def encode_token_weights(self, token_weight_pairs):
-        tok_pairs = token_weight_pairs["qwen25_7b"][0]
-        template_end = -1
-        if tok_pairs[0][0] == 27:
-            if len(tok_pairs) > 36:  # refiner prompt uses a fixed 36 template_end
-                template_end = 36
-
-        cond, p, extra = super().encode_token_weights(token_weight_pairs, template_end=template_end)
+        cond, p, extra = super().encode_token_weights(token_weight_pairs)
         if self.byt5_small is not None and "byt5" in token_weight_pairs:
             out = self.byt5_small.encode_token_weights(token_weight_pairs["byt5"])
             extra["conditioning_byt5small"] = out[0]

comfy/text_encoders/llama.py

@@ -3,7 +3,6 @@ import torch.nn as nn
 from dataclasses import dataclass
 from typing import Optional, Any
 import math
-import logging
 
 from comfy.ldm.modules.attention import optimized_attention_for_device
 import comfy.model_management
@@ -29,9 +28,6 @@ class Llama2Config:
     mlp_activation = "silu"
     qkv_bias = False
     rope_dims = None
-    q_norm = None
-    k_norm = None
-    rope_scale = None
 
 @dataclass
 class Qwen25_3BConfig:
@@ -50,9 +46,6 @@ class Qwen25_3BConfig:
     mlp_activation = "silu"
     qkv_bias = True
     rope_dims = None
-    q_norm = None
-    k_norm = None
-    rope_scale = None
 
 @dataclass
 class Qwen25_7BVLI_Config:
@@ -71,9 +64,6 @@ class Qwen25_7BVLI_Config:
     mlp_activation = "silu"
     qkv_bias = True
     rope_dims = [16, 24, 24]
-    q_norm = None
-    k_norm = None
-    rope_scale = None
 
 @dataclass
 class Gemma2_2B_Config:
@@ -92,32 +82,6 @@ class Gemma2_2B_Config:
     mlp_activation = "gelu_pytorch_tanh"
     qkv_bias = False
     rope_dims = None
-    q_norm = None
-    k_norm = None
-    sliding_attention = None
-    rope_scale = None
-
-@dataclass
-class Gemma3_4B_Config:
-    vocab_size: int = 262208
-    hidden_size: int = 2560
-    intermediate_size: int = 10240
-    num_hidden_layers: int = 34
-    num_attention_heads: int = 8
-    num_key_value_heads: int = 4
-    max_position_embeddings: int = 131072
-    rms_norm_eps: float = 1e-6
-    rope_theta = [10000.0, 1000000.0]
-    transformer_type: str = "gemma3"
-    head_dim = 256
-    rms_norm_add = True
-    mlp_activation = "gelu_pytorch_tanh"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    sliding_attention = [False, False, False, False, False, 1024]
-    rope_scale = [1.0, 8.0]
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5, add=False, device=None, dtype=None):
@@ -142,40 +106,25 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_dims=None, device=None):
-    if not isinstance(theta, list):
-        theta = [theta]
+def precompute_freqs_cis(head_dim, position_ids, theta, rope_dims=None, device=None):
+    theta_numerator = torch.arange(0, head_dim, 2, device=device).float()
+    inv_freq = 1.0 / (theta ** (theta_numerator / head_dim))
 
-    out = []
-    for index, t in enumerate(theta):
-        theta_numerator = torch.arange(0, head_dim, 2, device=device).float()
-        inv_freq = 1.0 / (t ** (theta_numerator / head_dim))
+    inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+    position_ids_expanded = position_ids[:, None, :].float()
+    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()
+    if rope_dims is not None and position_ids.shape[0] > 1:
+        mrope_section = rope_dims * 2
+        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0)
+        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0)
+    else:
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
 
-        if rope_scale is not None:
-            if isinstance(rope_scale, list):
-                inv_freq /= rope_scale[index]
-            else:
-                inv_freq /= rope_scale
-
-        inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-        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()
-        if rope_dims is not None and position_ids.shape[0] > 1:
-            mrope_section = rope_dims * 2
-            cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0)
-            sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0)
-        else:
-            cos = cos.unsqueeze(1)
-            sin = sin.unsqueeze(1)
-        out.append((cos, sin))
-
-    if len(out) == 1:
-        return out[0]
-
-    return out
+    return (cos, sin)
 
 
 def apply_rope(xq, xk, freqs_cis):
@@ -203,14 +152,6 @@ class Attention(nn.Module):
         self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
         self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype)
 
-        self.q_norm = None
-        self.k_norm = None
-
-        if config.q_norm == "gemma3":
-            self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
-        if config.k_norm == "gemma3":
-            self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -227,11 +168,6 @@ class Attention(nn.Module):
         xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
         xv = xv.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
-        if self.q_norm is not None:
-            xq = self.q_norm(xq)
-        if self.k_norm is not None:
-            xk = self.k_norm(xk)
-
         xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
 
         xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
@@ -256,7 +192,7 @@ class MLP(nn.Module):
         return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))
 
 class TransformerBlock(nn.Module):
-    def __init__(self, config: Llama2Config, index, device=None, dtype=None, ops: Any = None):
+    def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
         super().__init__()
         self.self_attn = Attention(config, device=device, dtype=dtype, ops=ops)
         self.mlp = MLP(config, device=device, dtype=dtype, ops=ops)
@@ -290,7 +226,7 @@ class TransformerBlock(nn.Module):
         return x
 
 class TransformerBlockGemma2(nn.Module):
-    def __init__(self, config: Llama2Config, index, device=None, dtype=None, ops: Any = None):
+    def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
         super().__init__()
         self.self_attn = Attention(config, device=device, dtype=dtype, ops=ops)
         self.mlp = MLP(config, device=device, dtype=dtype, ops=ops)
@@ -299,13 +235,6 @@ class TransformerBlockGemma2(nn.Module):
         self.pre_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
         self.post_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
 
-        if config.sliding_attention is not None:  # TODO: implement. (Not that necessary since models are trained on less than 1024 tokens)
-            self.sliding_attention = config.sliding_attention[index % len(config.sliding_attention)]
-        else:
-            self.sliding_attention = False
-
-        self.transformer_type = config.transformer_type
-
     def forward(
         self,
         x: torch.Tensor,
@@ -313,14 +242,6 @@ class TransformerBlockGemma2(nn.Module):
         freqs_cis: Optional[torch.Tensor] = None,
         optimized_attention=None,
     ):
-        if self.transformer_type == 'gemma3':
-            if self.sliding_attention:
-                if x.shape[1] > self.sliding_attention:
-                    logging.warning("Warning: sliding attention not implemented, results may be incorrect")
-                freqs_cis = freqs_cis[1]
-            else:
-                freqs_cis = freqs_cis[0]
-
         # Self Attention
         residual = x
         x = self.input_layernorm(x)
@@ -355,7 +276,7 @@ class Llama2_(nn.Module):
             device=device,
             dtype=dtype
         )
-        if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
+        if self.config.transformer_type == "gemma2":
             transformer = TransformerBlockGemma2
             self.normalize_in = True
         else:
@@ -363,8 +284,8 @@ class Llama2_(nn.Module):
             self.normalize_in = False
 
         self.layers = nn.ModuleList([
-            transformer(config, index=i, device=device, dtype=dtype, ops=ops)
-            for i in range(config.num_hidden_layers)
+            transformer(config, device=device, dtype=dtype, ops=ops)
+            for _ in range(config.num_hidden_layers)
         ])
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
         # self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
@@ -384,7 +305,6 @@ class Llama2_(nn.Module):
         freqs_cis = precompute_freqs_cis(self.config.head_dim,
                                          position_ids,
                                          self.config.rope_theta,
-                                         self.config.rope_scale,
                                          self.config.rope_dims,
                                          device=x.device)
 
@@ -513,12 +433,3 @@ class Gemma2_2B(BaseLlama, torch.nn.Module):
 
         self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
         self.dtype = dtype
-
-class Gemma3_4B(BaseLlama, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Gemma3_4B_Config(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype

comfy/text_encoders/lumina2.py

@@ -11,41 +11,23 @@ class Gemma2BTokenizer(sd1_clip.SDTokenizer):
     def state_dict(self):
         return {"spiece_model": self.tokenizer.serialize_model()}
 
-class Gemma3_4BTokenizer(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=2560, embedding_key='gemma3_4b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, tokenizer_data=tokenizer_data)
-
-    def state_dict(self):
-        return {"spiece_model": self.tokenizer.serialize_model()}
 
 class LuminaTokenizer(sd1_clip.SD1Tokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma2_2b", tokenizer=Gemma2BTokenizer)
 
-class NTokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma3_4b", tokenizer=Gemma3_4BTokenizer)
 
 class Gemma2_2BModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}):
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma2_2B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-class Gemma3_4BModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}):
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
 class LuminaModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}, name="gemma2_2b", clip_model=Gemma2_2BModel):
-        super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        super().__init__(device=device, dtype=dtype, name="gemma2_2b", clip_model=Gemma2_2BModel, model_options=model_options)
 
 
-def te(dtype_llama=None, llama_scaled_fp8=None, model_type="gemma2_2b"):
-    if model_type == "gemma2_2b":
-        model = Gemma2_2BModel
-    elif model_type == "gemma3_4b":
-        model = Gemma3_4BModel
-
+def te(dtype_llama=None, llama_scaled_fp8=None):
     class LuminaTEModel_(LuminaModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
@@ -53,5 +35,5 @@ def te(dtype_llama=None, llama_scaled_fp8=None, model_type="gemma2_2b"):
                 model_options["scaled_fp8"] = llama_scaled_fp8
             if dtype_llama is not None:
                 dtype = dtype_llama
-            super().__init__(device=device, dtype=dtype, name=model_type, model_options=model_options, clip_model=model)
+            super().__init__(device=device, dtype=dtype, model_options=model_options)
     return LuminaTEModel_

comfy/text_encoders/qwen_image.py

@@ -18,22 +18,13 @@ class QwenImageTokenizer(sd1_clip.SD1Tokenizer):
         self.llama_template_images = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
 
     def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], **kwargs):
-        skip_template = False
-        if text.startswith('<|im_start|>'):
-            skip_template = True
-        if text.startswith('<|start_header_id|>'):
-            skip_template = True
-
-        if skip_template:
-            llama_text = text
-        else:
-            if llama_template is None:
-                if len(images) > 0:
-                    llama_text = self.llama_template_images.format(text)
-                else:
-                    llama_text = self.llama_template.format(text)
+        if llama_template is None:
+            if len(images) > 0:
+                llama_text = self.llama_template_images.format(text)
             else:
-                llama_text = llama_template.format(text)
+                llama_text = self.llama_template.format(text)
+        else:
+            llama_text = llama_template.format(text)
         tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
         key_name = next(iter(tokens))
         embed_count = 0
@@ -56,23 +47,22 @@ class QwenImageTEModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}):
         super().__init__(device=device, dtype=dtype, name="qwen25_7b", clip_model=Qwen25_7BVLIModel, model_options=model_options)
 
-    def encode_token_weights(self, token_weight_pairs, template_end=-1):
+    def encode_token_weights(self, token_weight_pairs):
         out, pooled, extra = super().encode_token_weights(token_weight_pairs)
         tok_pairs = token_weight_pairs["qwen25_7b"][0]
         count_im_start = 0
-        if template_end == -1:
-            for i, v in enumerate(tok_pairs):
-                elem = v[0]
-                if not torch.is_tensor(elem):
-                    if isinstance(elem, numbers.Integral):
-                        if elem == 151644 and count_im_start < 2:
-                            template_end = i
-                            count_im_start += 1
+        for i, v in enumerate(tok_pairs):
+            elem = v[0]
+            if not torch.is_tensor(elem):
+                if isinstance(elem, numbers.Integral):
+                    if elem == 151644 and count_im_start < 2:
+                        template_end = i
+                        count_im_start += 1
 
-            if out.shape[1] > (template_end + 3):
-                if tok_pairs[template_end + 1][0] == 872:
-                    if tok_pairs[template_end + 2][0] == 198:
-                        template_end += 3
+        if out.shape[1] > (template_end + 3):
+            if tok_pairs[template_end + 1][0] == 872:
+                if tok_pairs[template_end + 2][0] == 198:
+                    template_end += 3
 
         out = out[:, template_end:]
 

comfy_api/latest/__init__.py

@@ -8,8 +8,8 @@ from comfy_api.internal.async_to_sync import create_sync_class
 from comfy_api.latest._input import ImageInput, AudioInput, MaskInput, LatentInput, VideoInput
 from comfy_api.latest._input_impl import VideoFromFile, VideoFromComponents
 from comfy_api.latest._util import VideoCodec, VideoContainer, VideoComponents
-from . import _io as io
-from . import _ui as ui
+from comfy_api.latest._io import _IO as io  #noqa: F401
+from comfy_api.latest._ui import _UI as ui  #noqa: F401
 # from comfy_api.latest._resources import _RESOURCES as resources  #noqa: F401
 from comfy_execution.utils import get_executing_context
 from comfy_execution.progress import get_progress_state, PreviewImageTuple
@@ -114,8 +114,6 @@ if TYPE_CHECKING:
     ComfyAPISync: Type[comfy_api.latest.generated.ComfyAPISyncStub.ComfyAPISyncStub]
 ComfyAPISync = create_sync_class(ComfyAPI_latest)
 
-comfy_io = io  # create the new alias for io
-
 __all__ = [
     "ComfyAPI",
     "ComfyAPISync",
@@ -123,7 +121,4 @@ __all__ = [
     "InputImpl",
     "Types",
     "ComfyExtension",
-    "io",
-    "comfy_io",
-    "ui",
 ]

comfy_api/latest/_input/video_types.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
-from typing import Optional, Union, IO
+from typing import Optional, Union
 import io
 import av
 from comfy_api.util import VideoContainer, VideoCodec, VideoComponents
@@ -23,7 +23,7 @@ class VideoInput(ABC):
     @abstractmethod
     def save_to(
         self,
-        path: Union[str, IO[bytes]],
+        path: str,
         format: VideoContainer = VideoContainer.AUTO,
         codec: VideoCodec = VideoCodec.AUTO,
         metadata: Optional[dict] = None

comfy_api/latest/_io.py

@@ -336,25 +336,11 @@ class Combo(ComfyTypeIO):
     class Input(WidgetInput):
         """Combo input (dropdown)."""
         Type = str
-        def __init__(
-            self,
-            id: str,
-            options: list[str] | list[int] | type[Enum] = None,
-            display_name: str=None,
-            optional=False,
-            tooltip: str=None,
-            lazy: bool=None,
-            default: str | int | Enum = None,
-            control_after_generate: bool=None,
-            upload: UploadType=None,
-            image_folder: FolderType=None,
-            remote: RemoteOptions=None,
-            socketless: bool=None,
-        ):
-            if isinstance(options, type) and issubclass(options, Enum):
-                options = [v.value for v in options]
-            if isinstance(default, Enum):
-                default = default.value
+        def __init__(self, id: str, options: list[str]=None, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
+                    default: str=None, control_after_generate: bool=None,
+                    upload: UploadType=None, image_folder: FolderType=None,
+                    remote: RemoteOptions=None,
+                    socketless: bool=None):
             super().__init__(id, display_name, optional, tooltip, lazy, default, socketless)
             self.multiselect = False
             self.options = options
@@ -1582,78 +1568,77 @@ class _UIOutput(ABC):
         ...
 
 
-__all__ = [
-    "FolderType",
-    "UploadType",
-    "RemoteOptions",
-    "NumberDisplay",
+class _IO:
+    FolderType = FolderType
+    UploadType = UploadType
+    RemoteOptions = RemoteOptions
+    NumberDisplay = NumberDisplay
 
-    "comfytype",
-    "Custom",
-    "Input",
-    "WidgetInput",
-    "Output",
-    "ComfyTypeI",
-    "ComfyTypeIO",
+    comfytype = staticmethod(comfytype)
+    Custom = staticmethod(Custom)
+    Input = Input
+    WidgetInput = WidgetInput
+    Output = Output
+    ComfyTypeI = ComfyTypeI
+    ComfyTypeIO = ComfyTypeIO
+    #---------------------------------
     # Supported Types
-    "Boolean",
-    "Int",
-    "Float",
-    "String",
-    "Combo",
-    "MultiCombo",
-    "Image",
-    "WanCameraEmbedding",
-    "Webcam",
-    "Mask",
-    "Latent",
-    "Conditioning",
-    "Sampler",
-    "Sigmas",
-    "Noise",
-    "Guider",
-    "Clip",
-    "ControlNet",
-    "Vae",
-    "Model",
-    "ClipVision",
-    "ClipVisionOutput",
-    "AudioEncoder",
-    "AudioEncoderOutput",
-    "StyleModel",
-    "Gligen",
-    "UpscaleModel",
-    "Audio",
-    "Video",
-    "SVG",
-    "LoraModel",
-    "LossMap",
-    "Voxel",
-    "Mesh",
-    "Hooks",
-    "HookKeyframes",
-    "TimestepsRange",
-    "LatentOperation",
-    "FlowControl",
-    "Accumulation",
-    "Load3DCamera",
-    "Load3D",
-    "Load3DAnimation",
-    "Photomaker",
-    "Point",
-    "FaceAnalysis",
-    "BBOX",
-    "SEGS",
-    "AnyType",
-    "MultiType",
-    # Other classes
-    "HiddenHolder",
-    "Hidden",
-    "NodeInfoV1",
-    "NodeInfoV3",
-    "Schema",
-    "ComfyNode",
-    "NodeOutput",
-    "add_to_dict_v1",
-    "add_to_dict_v3",
-]
+    Boolean = Boolean
+    Int = Int
+    Float = Float
+    String = String
+    Combo = Combo
+    MultiCombo = MultiCombo
+    Image = Image
+    WanCameraEmbedding = WanCameraEmbedding
+    Webcam = Webcam
+    Mask = Mask
+    Latent = Latent
+    Conditioning = Conditioning
+    Sampler = Sampler
+    Sigmas = Sigmas
+    Noise = Noise
+    Guider = Guider
+    Clip = Clip
+    ControlNet = ControlNet
+    Vae = Vae
+    Model = Model
+    ClipVision = ClipVision
+    ClipVisionOutput = ClipVisionOutput
+    AudioEncoderOutput = AudioEncoderOutput
+    StyleModel = StyleModel
+    Gligen = Gligen
+    UpscaleModel = UpscaleModel
+    Audio = Audio
+    Video = Video
+    SVG = SVG
+    LoraModel = LoraModel
+    LossMap = LossMap
+    Voxel = Voxel
+    Mesh = Mesh
+    Hooks = Hooks
+    HookKeyframes = HookKeyframes
+    TimestepsRange = TimestepsRange
+    LatentOperation = LatentOperation
+    FlowControl = FlowControl
+    Accumulation = Accumulation
+    Load3DCamera = Load3DCamera
+    Load3D = Load3D
+    Load3DAnimation = Load3DAnimation
+    Photomaker = Photomaker
+    Point = Point
+    FaceAnalysis = FaceAnalysis
+    BBOX = BBOX
+    SEGS = SEGS
+    AnyType = AnyType
+    MultiType = MultiType
+    #---------------------------------
+    HiddenHolder = HiddenHolder
+    Hidden = Hidden
+    NodeInfoV1 = NodeInfoV1
+    NodeInfoV3 = NodeInfoV3
+    Schema = Schema
+    ComfyNode = ComfyNode
+    NodeOutput = NodeOutput
+    add_to_dict_v1 = staticmethod(add_to_dict_v1)
+    add_to_dict_v3 = staticmethod(add_to_dict_v3)

comfy_api/latest/_ui.py

@@ -449,16 +449,15 @@ class PreviewText(_UIOutput):
         return {"text": (self.value,)}
 
 
-__all__ = [
-    "SavedResult",
-    "SavedImages",
-    "SavedAudios",
-    "ImageSaveHelper",
-    "AudioSaveHelper",
-    "PreviewImage",
-    "PreviewMask",
-    "PreviewAudio",
-    "PreviewVideo",
-    "PreviewUI3D",
-    "PreviewText",
-]
+class _UI:
+    SavedResult = SavedResult
+    SavedImages = SavedImages
+    SavedAudios = SavedAudios
+    ImageSaveHelper = ImageSaveHelper
+    AudioSaveHelper = AudioSaveHelper
+    PreviewImage = PreviewImage
+    PreviewMask = PreviewMask
+    PreviewAudio = PreviewAudio
+    PreviewVideo = PreviewVideo
+    PreviewUI3D = PreviewUI3D
+    PreviewText = PreviewText

comfy_api_nodes/apinode_utils.py

@@ -18,7 +18,7 @@ from comfy_api_nodes.apis.client import (
     UploadResponse,
 )
 from server import PromptServer
-from comfy.cli_args import args
+
 
 import numpy as np
 from PIL import Image
@@ -30,9 +30,7 @@ from io import BytesIO
 import av
 
 
-async def download_url_to_video_output(
-    video_url: str, timeout: int = None, auth_kwargs: Optional[dict[str, str]] = None
-) -> VideoFromFile:
+async def download_url_to_video_output(video_url: str, timeout: int = None) -> VideoFromFile:
     """Downloads a video from a URL and returns a `VIDEO` output.
 
     Args:
@@ -41,7 +39,7 @@ async def download_url_to_video_output(
     Returns:
         A Comfy node `VIDEO` output.
     """
-    video_io = await download_url_to_bytesio(video_url, timeout, auth_kwargs=auth_kwargs)
+    video_io = await download_url_to_bytesio(video_url, timeout)
     if video_io is None:
         error_msg = f"Failed to download video from {video_url}"
         logging.error(error_msg)
@@ -154,7 +152,7 @@ def validate_aspect_ratio(
             raise TypeError(
                 f"Aspect ratio cannot reduce to any less than {minimum_ratio_str} ({minimum_ratio}), but was {aspect_ratio} ({calculated_ratio})."
             )
-        if calculated_ratio > maximum_ratio:
+        elif calculated_ratio > maximum_ratio:
             raise TypeError(
                 f"Aspect ratio cannot reduce to any greater than {maximum_ratio_str} ({maximum_ratio}), but was {aspect_ratio} ({calculated_ratio})."
             )
@@ -166,9 +164,7 @@ def mimetype_to_extension(mime_type: str) -> str:
     return mime_type.split("/")[-1].lower()
 
 
-async def download_url_to_bytesio(
-    url: str, timeout: int = None, auth_kwargs: Optional[dict[str, str]] = None
-) -> BytesIO:
+async def download_url_to_bytesio(url: str, timeout: int = None) -> BytesIO:
     """Downloads content from a URL using requests and returns it as BytesIO.
 
     Args:
@@ -178,18 +174,9 @@ async def download_url_to_bytesio(
     Returns:
         BytesIO object containing the downloaded content.
     """
-    headers = {}
-    if url.startswith("/proxy/"):
-        url = str(args.comfy_api_base).rstrip("/") + url
-        auth_token = auth_kwargs.get("auth_token")
-        comfy_api_key = auth_kwargs.get("comfy_api_key")
-        if auth_token:
-            headers["Authorization"] = f"Bearer {auth_token}"
-        elif comfy_api_key:
-            headers["X-API-KEY"] = comfy_api_key
     timeout_cfg = aiohttp.ClientTimeout(total=timeout) if timeout else None
     async with aiohttp.ClientSession(timeout=timeout_cfg) as session:
-        async with session.get(url, headers=headers) as resp:
+        async with session.get(url) as resp:
             resp.raise_for_status()  # Raises HTTPError for bad responses (4XX or 5XX)
             return BytesIO(await resp.read())
 
@@ -269,7 +256,7 @@ def tensor_to_bytesio(
         mime_type: Target image MIME type (e.g., 'image/png', 'image/jpeg', 'image/webp', 'video/mp4').
 
     Returns:
-        Named BytesIO object containing the image data, with pointer set to the start of buffer.
+        Named BytesIO object containing the image data.
     """
     if not mime_type:
         mime_type = "image/png"
@@ -431,7 +418,7 @@ async def upload_video_to_comfyapi(
                     f"Video duration ({actual_duration:.2f}s) exceeds the maximum allowed ({max_duration}s)."
                 )
         except Exception as e:
-            logging.error("Error getting video duration: %s", str(e))
+            logging.error(f"Error getting video duration: {e}")
             raise ValueError(f"Could not verify video duration from source: {e}") from e
 
     upload_mime_type = f"video/{container.value.lower()}"

comfy_api_nodes/apis/__init__.py

@@ -2,7 +2,6 @@
 #   filename:  filtered-openapi.yaml
 #   timestamp: 2025-07-30T08:54:00+00:00
 
-# pylint: disable
 from __future__ import annotations
 
 from datetime import date, datetime
@@ -1321,7 +1320,6 @@ class KlingTextToVideoModelName(str, Enum):
     kling_v1 = 'kling-v1'
     kling_v1_6 = 'kling-v1-6'
     kling_v2_1_master = 'kling-v2-1-master'
-    kling_v2_5_turbo = 'kling-v2-5-turbo'
 
 
 class KlingVideoGenAspectRatio(str, Enum):
@@ -1356,7 +1354,6 @@ class KlingVideoGenModelName(str, Enum):
     kling_v2_master = 'kling-v2-master'
     kling_v2_1 = 'kling-v2-1'
     kling_v2_1_master = 'kling-v2-1-master'
-    kling_v2_5_turbo = 'kling-v2-5-turbo'
 
 
 class KlingVideoResult(BaseModel):

comfy_api_nodes/apis/client.py

@@ -95,10 +95,9 @@ import aiohttp
 import asyncio
 import logging
 import io
-import os
 import socket
 from aiohttp.client_exceptions import ClientError, ClientResponseError
-from typing import Type, Optional, Any, TypeVar, Generic, Callable
+from typing import Dict, Type, Optional, Any, TypeVar, Generic, Callable, Tuple
 from enum import Enum
 import json
 from urllib.parse import urljoin, urlparse
@@ -175,7 +174,7 @@ class ApiClient:
         max_retries: int = 3,
         retry_delay: float = 1.0,
         retry_backoff_factor: float = 2.0,
-        retry_status_codes: Optional[tuple[int, ...]] = None,
+        retry_status_codes: Optional[Tuple[int, ...]] = None,
         session: Optional[aiohttp.ClientSession] = None,
     ):
         self.base_url = base_url
@@ -199,9 +198,9 @@ class ApiClient:
 
     @staticmethod
     def _create_json_payload_args(
-        data: Optional[dict[str, Any]] = None,
-        headers: Optional[dict[str, str]] = None,
-    ) -> dict[str, Any]:
+        data: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
         return {
             "json": data,
             "headers": headers,
@@ -209,27 +208,24 @@ class ApiClient:
 
     def _create_form_data_args(
         self,
-        data: dict[str, Any] | None,
-        files: dict[str, Any] | None,
-        headers: Optional[dict[str, str]] = None,
+        data: Dict[str, Any] | None,
+        files: Dict[str, Any] | None,
+        headers: Optional[Dict[str, str]] = None,
         multipart_parser: Callable | None = None,
-    ) -> dict[str, Any]:
+    ) -> Dict[str, Any]:
         if headers and "Content-Type" in headers:
             del headers["Content-Type"]
 
         if multipart_parser and data:
             data = multipart_parser(data)
 
-        if isinstance(data, aiohttp.FormData):
-            form = data  # If the parser already returned a FormData, pass it through
-        else:
-            form = aiohttp.FormData(default_to_multipart=True)
-            if data:  # regular text fields
-                for k, v in data.items():
-                    if v is None:
-                        continue  # aiohttp fails to serialize "None" values
-                    # aiohttp expects strings or bytes; convert enums etc.
-                    form.add_field(k, str(v) if not isinstance(v, (bytes, bytearray)) else v)
+        form = aiohttp.FormData(default_to_multipart=True)
+        if data:  # regular text fields
+            for k, v in data.items():
+                if v is None:
+                    continue  # aiohttp fails to serialize "None" values
+                # aiohttp expects strings or bytes; convert enums etc.
+                form.add_field(k, str(v) if not isinstance(v, (bytes, bytearray)) else v)
 
         if files:
             file_iter = files if isinstance(files, list) else files.items()
@@ -254,9 +250,9 @@ class ApiClient:
 
     @staticmethod
     def _create_urlencoded_form_data_args(
-        data: dict[str, Any],
-        headers: Optional[dict[str, str]] = None,
-    ) -> dict[str, Any]:
+        data: Dict[str, Any],
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
         headers = headers or {}
         headers["Content-Type"] = "application/x-www-form-urlencoded"
         return {
@@ -264,7 +260,7 @@ class ApiClient:
             "headers": headers,
         }
 
-    def get_headers(self) -> dict[str, str]:
+    def get_headers(self) -> Dict[str, str]:
         """Get headers for API requests, including authentication if available"""
         headers = {"Content-Type": "application/json", "Accept": "application/json"}
 
@@ -275,7 +271,7 @@ class ApiClient:
 
         return headers
 
-    async def _check_connectivity(self, target_url: str) -> dict[str, bool]:
+    async def _check_connectivity(self, target_url: str) -> Dict[str, bool]:
         """
         Check connectivity to determine if network issues are local or server-related.
 
@@ -316,14 +312,14 @@ class ApiClient:
         self,
         method: str,
         path: str,
-        params: Optional[dict[str, Any]] = None,
-        data: Optional[dict[str, Any]] = None,
-        files: Optional[dict[str, Any] | list[tuple[str, Any]]] = None,
-        headers: Optional[dict[str, str]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
+        files: Optional[Dict[str, Any] | list[tuple[str, Any]]] = None,
+        headers: Optional[Dict[str, str]] = None,
         content_type: str = "application/json",
         multipart_parser: Callable | None = None,
         retry_count: int = 0,  # Used internally for tracking retries
-    ) -> dict[str, Any]:
+    ) -> Dict[str, Any]:
         """
         Make an HTTP request to the API with automatic retries for transient errors.
 
@@ -359,10 +355,10 @@ class ApiClient:
         if params:
             params = {k: v for k, v in params.items() if v is not None}  # aiohttp fails to serialize None values
 
-        logging.debug("[DEBUG] Request Headers: %s", request_headers)
-        logging.debug("[DEBUG] Files: %s", files)
-        logging.debug("[DEBUG] Params: %s", params)
-        logging.debug("[DEBUG] Data: %s", data)
+        logging.debug(f"[DEBUG] Request Headers: {request_headers}")
+        logging.debug(f"[DEBUG] Files: {files}")
+        logging.debug(f"[DEBUG] Params: {params}")
+        logging.debug(f"[DEBUG] Data: {data}")
 
         if content_type == "application/x-www-form-urlencoded":
             payload_args = self._create_urlencoded_form_data_args(data or {}, request_headers)
@@ -485,7 +481,7 @@ class ApiClient:
             retry_delay: Initial delay between retries in seconds
             retry_backoff_factor: Multiplier for the delay after each retry
         """
-        headers: dict[str, str] = {}
+        headers: Dict[str, str] = {}
         skip_auto_headers: set[str] = set()
         if content_type:
             headers["Content-Type"] = content_type
@@ -503,9 +499,7 @@ class ApiClient:
         else:
             raise ValueError("File must be BytesIO or str path")
 
-        parsed = urlparse(upload_url)
-        basename = os.path.basename(parsed.path) or parsed.netloc or "upload"
-        operation_id = f"upload_{basename}_{uuid.uuid4().hex[:8]}"
+        operation_id = f"upload_{upload_url.split('/')[-1]}_{uuid.uuid4().hex[:8]}"
         request_logger.log_request_response(
             operation_id=operation_id,
             request_method="PUT",
@@ -538,7 +532,7 @@ class ApiClient:
                     request_method="PUT",
                     request_url=upload_url,
                     response_status_code=e.status if hasattr(e, "status") else None,
-                    response_headers=dict(e.headers) if hasattr(e, "headers") else None,
+                    response_headers=dict(e.headers) if getattr(e, "headers") else None,
                     response_content=None,
                     error_message=f"{type(e).__name__}: {str(e)}",
                 )
@@ -558,7 +552,7 @@ class ApiClient:
         *req_meta,
         retry_count: int,
         response_content: dict | str = "",
-    ) -> dict[str, Any]:
+    ) -> Dict[str, Any]:
         status_code = exc.status
         if status_code == 401:
             user_friendly = "Unauthorized: Please login first to use this node."
@@ -592,9 +586,9 @@ class ApiClient:
             error_message=f"HTTP Error {exc.status}",
         )
 
-        logging.debug("[DEBUG] API Error: %s (Status: %s)", user_friendly, status_code)
+        logging.debug(f"[DEBUG] API Error: {user_friendly} (Status: {status_code})")
         if response_content:
-            logging.debug("[DEBUG] Response content: %s", response_content)
+            logging.debug(f"[DEBUG] Response content: {response_content}")
 
         # Retry if eligible
         if status_code in self.retry_status_codes and retry_count < self.max_retries:
@@ -659,7 +653,7 @@ class ApiEndpoint(Generic[T, R]):
         method: HttpMethod,
         request_model: Type[T],
         response_model: Type[R],
-        query_params: Optional[dict[str, Any]] = None,
+        query_params: Optional[Dict[str, Any]] = None,
     ):
         """Initialize an API endpoint definition.
 
@@ -684,11 +678,11 @@ class SynchronousOperation(Generic[T, R]):
         self,
         endpoint: ApiEndpoint[T, R],
         request: T,
-        files: Optional[dict[str, Any] | list[tuple[str, Any]]] = None,
+        files: Optional[Dict[str, Any] | list[tuple[str, Any]]] = None,
         api_base: str | None = None,
         auth_token: Optional[str] = None,
         comfy_api_key: Optional[str] = None,
-        auth_kwargs: Optional[dict[str, str]] = None,
+        auth_kwargs: Optional[Dict[str, str]] = None,
         timeout: float = 7200.0,
         verify_ssl: bool = True,
         content_type: str = "application/json",
@@ -729,7 +723,7 @@ class SynchronousOperation(Generic[T, R]):
             )
 
         try:
-            request_dict: Optional[dict[str, Any]]
+            request_dict: Optional[Dict[str, Any]]
             if isinstance(self.request, EmptyRequest):
                 request_dict = None
             else:
@@ -738,9 +732,11 @@ class SynchronousOperation(Generic[T, R]):
                     if isinstance(v, Enum):
                         request_dict[k] = v.value
 
-            logging.debug("[DEBUG] API Request: %s %s", self.endpoint.method.value, self.endpoint.path)
-            logging.debug("[DEBUG] Request Data: %s", json.dumps(request_dict, indent=2))
-            logging.debug("[DEBUG] Query Params: %s", self.endpoint.query_params)
+            logging.debug(
+                f"[DEBUG] API Request: {self.endpoint.method.value} {self.endpoint.path}"
+            )
+            logging.debug(f"[DEBUG] Request Data: {json.dumps(request_dict, indent=2)}")
+            logging.debug(f"[DEBUG] Query Params: {self.endpoint.query_params}")
 
             response_json = await client.request(
                 self.endpoint.method.value,
@@ -755,11 +751,11 @@ class SynchronousOperation(Generic[T, R]):
             logging.debug("=" * 50)
             logging.debug("[DEBUG] RESPONSE DETAILS:")
             logging.debug("[DEBUG] Status Code: 200 (Success)")
-            logging.debug("[DEBUG] Response Body: %s", json.dumps(response_json, indent=2))
+            logging.debug(f"[DEBUG] Response Body: {json.dumps(response_json, indent=2)}")
             logging.debug("=" * 50)
 
             parsed_response = self.endpoint.response_model.model_validate(response_json)
-            logging.debug("[DEBUG] Parsed Response: %s", parsed_response)
+            logging.debug(f"[DEBUG] Parsed Response: {parsed_response}")
             return parsed_response
         finally:
             if owns_client:
@@ -782,14 +778,14 @@ class PollingOperation(Generic[T, R]):
         poll_endpoint: ApiEndpoint[EmptyRequest, R],
         completed_statuses: list[str],
         failed_statuses: list[str],
-        status_extractor: Callable[[R], Optional[str]],
-        progress_extractor: Callable[[R], Optional[float]] | None = None,
-        result_url_extractor: Callable[[R], Optional[str]] | None = None,
+        status_extractor: Callable[[R], str],
+        progress_extractor: Callable[[R], float] | None = None,
+        result_url_extractor: Callable[[R], str] | None = None,
         request: Optional[T] = None,
         api_base: str | None = None,
         auth_token: Optional[str] = None,
         comfy_api_key: Optional[str] = None,
-        auth_kwargs: Optional[dict[str, str]] = None,
+        auth_kwargs: Optional[Dict[str, str]] = None,
         poll_interval: float = 5.0,
         max_poll_attempts: int = 120,  # Default max polling attempts (10 minutes with 5s interval)
         max_retries: int = 3,  # Max retries per individual API call
@@ -875,7 +871,7 @@ class PollingOperation(Generic[T, R]):
         status = TaskStatus.PENDING
         for poll_count in range(1, self.max_poll_attempts + 1):
             try:
-                logging.debug("[DEBUG] Polling attempt #%s", poll_count)
+                logging.debug(f"[DEBUG] Polling attempt #{poll_count}")
 
                 request_dict = (
                     None if self.request is None else self.request.model_dump(exclude_none=True)
@@ -883,13 +879,10 @@ class PollingOperation(Generic[T, R]):
 
                 if poll_count == 1:
                     logging.debug(
-                        "[DEBUG] Poll Request: %s %s",
-                        self.poll_endpoint.method.value,
-                        self.poll_endpoint.path,
+                        f"[DEBUG] Poll Request: {self.poll_endpoint.method.value} {self.poll_endpoint.path}"
                     )
                     logging.debug(
-                        "[DEBUG] Poll Request Data: %s",
-                        json.dumps(request_dict, indent=2) if request_dict else "None",
+                        f"[DEBUG] Poll Request Data: {json.dumps(request_dict, indent=2) if request_dict else 'None'}"
                     )
 
                 # Query task status
@@ -904,7 +897,7 @@ class PollingOperation(Generic[T, R]):
 
                 # Check if task is complete
                 status = self._check_task_status(response_obj)
-                logging.debug("[DEBUG] Task Status: %s", status)
+                logging.debug(f"[DEBUG] Task Status: {status}")
 
                 # If progress extractor is provided, extract progress
                 if self.progress_extractor:
@@ -918,7 +911,7 @@ class PollingOperation(Generic[T, R]):
                         result_url = self.result_url_extractor(response_obj)
                         if result_url:
                             message = f"Result URL: {result_url}"
-                    logging.debug("[DEBUG] %s", message)
+                    logging.debug(f"[DEBUG] {message}")
                     self._display_text_on_node(message)
                     self.final_response = response_obj
                     if self.progress_extractor:
@@ -926,7 +919,7 @@ class PollingOperation(Generic[T, R]):
                     return self.final_response
                 if status == TaskStatus.FAILED:
                     message = f"Task failed: {json.dumps(resp)}"
-                    logging.error("[DEBUG] %s", message)
+                    logging.error(f"[DEBUG] {message}")
                     raise Exception(message)
                 logging.debug("[DEBUG] Task still pending, continuing to poll...")
                 # Task pending – wait
@@ -940,12 +933,7 @@ class PollingOperation(Generic[T, R]):
                     raise Exception(
                         f"Polling aborted after {consecutive_errors} network errors: {str(e)}"
                     ) from e
-                logging.warning(
-                    "Network error (%s/%s): %s",
-                    consecutive_errors,
-                    max_consecutive_errors,
-                    str(e),
-                )
+                logging.warning("Network error (%s/%s): %s", consecutive_errors, max_consecutive_errors, str(e))
                 await asyncio.sleep(self.poll_interval)
             except Exception as e:
                 # For other errors, increment count and potentially abort
@@ -955,13 +943,10 @@ class PollingOperation(Generic[T, R]):
                         f"Polling aborted after {consecutive_errors} consecutive errors: {str(e)}"
                     ) from e
 
-                logging.error("[DEBUG] Polling error: %s", str(e))
+                logging.error(f"[DEBUG] Polling error: {str(e)}")
                 logging.warning(
-                    "Error during polling (attempt %s/%s): %s. Will retry in %s seconds.",
-                    poll_count,
-                    self.max_poll_attempts,
-                    str(e),
-                    self.poll_interval,
+                    f"Error during polling (attempt {poll_count}/{self.max_poll_attempts}): {str(e)}. "
+                    f"Will retry in {self.poll_interval} seconds."
                 )
                 await asyncio.sleep(self.poll_interval)
 

comfy_api_nodes/apis/pika_defs.py

@@ -1,100 +0,0 @@
-from typing import Optional
-from enum import Enum
-from pydantic import BaseModel, Field
-
-
-class Pikaffect(str, Enum):
-    Cake_ify = "Cake-ify"
-    Crumble = "Crumble"
-    Crush = "Crush"
-    Decapitate = "Decapitate"
-    Deflate = "Deflate"
-    Dissolve = "Dissolve"
-    Explode = "Explode"
-    Eye_pop = "Eye-pop"
-    Inflate = "Inflate"
-    Levitate = "Levitate"
-    Melt = "Melt"
-    Peel = "Peel"
-    Poke = "Poke"
-    Squish = "Squish"
-    Ta_da = "Ta-da"
-    Tear = "Tear"
-
-
-class PikaBodyGenerate22C2vGenerate22PikascenesPost(BaseModel):
-    aspectRatio: Optional[float] = Field(None, description='Aspect ratio (width / height)')
-    duration: Optional[int] = Field(5)
-    ingredientsMode: str = Field(...)
-    negativePrompt: Optional[str] = Field(None)
-    promptText: Optional[str] = Field(None)
-    resolution: Optional[str] = Field('1080p')
-    seed: Optional[int] = Field(None)
-
-
-class PikaGenerateResponse(BaseModel):
-    video_id: str = Field(...)
-
-
-class PikaBodyGenerate22I2vGenerate22I2vPost(BaseModel):
-    duration: Optional[int] = 5
-    negativePrompt: Optional[str] = Field(None)
-    promptText: Optional[str] = Field(None)
-    resolution: Optional[str] = '1080p'
-    seed: Optional[int] = Field(None)
-
-
-class PikaBodyGenerate22KeyframeGenerate22PikaframesPost(BaseModel):
-    duration: Optional[int] = Field(None, ge=5, le=10)
-    negativePrompt: Optional[str] = Field(None)
-    promptText: str = Field(...)
-    resolution: Optional[str] = '1080p'
-    seed: Optional[int] = Field(None)
-
-
-class PikaBodyGenerate22T2vGenerate22T2vPost(BaseModel):
-    aspectRatio: Optional[float] = Field(
-        1.7777777777777777,
-        description='Aspect ratio (width / height)',
-        ge=0.4,
-        le=2.5,
-    )
-    duration: Optional[int] = 5
-    negativePrompt: Optional[str] = Field(None)
-    promptText: str = Field(...)
-    resolution: Optional[str] = '1080p'
-    seed: Optional[int] = Field(None)
-
-
-class PikaBodyGeneratePikadditionsGeneratePikadditionsPost(BaseModel):
-    negativePrompt: Optional[str] = Field(None)
-    promptText: Optional[str] = Field(None)
-    seed: Optional[int] = Field(None)
-
-
-class PikaBodyGeneratePikaffectsGeneratePikaffectsPost(BaseModel):
-    negativePrompt: Optional[str] = Field(None)
-    pikaffect: Optional[str] = None
-    promptText: Optional[str] = Field(None)
-    seed: Optional[int] = Field(None)
-
-
-class PikaBodyGeneratePikaswapsGeneratePikaswapsPost(BaseModel):
-    negativePrompt: Optional[str] = Field(None)
-    promptText: Optional[str] = Field(None)
-    seed: Optional[int] = Field(None)
-    modifyRegionRoi: Optional[str] = Field(None)
-
-
-class PikaStatusEnum(str, Enum):
-    queued = "queued"
-    started = "started"
-    finished = "finished"
-    failed = "failed"
-
-
-class PikaVideoResponse(BaseModel):
-    id: str = Field(...)
-    progress: Optional[int] = Field(None)
-    status: PikaStatusEnum
-    url: Optional[str] = Field(None)

comfy_api_nodes/apis/request_logger.py

@@ -4,99 +4,62 @@ import os
 import datetime
 import json
 import logging
-import re
-import hashlib
-from typing import Any
-
 import folder_paths
 
 # Get the logger instance
 logger = logging.getLogger(__name__)
 
-
 def get_log_directory():
-    """Ensures the API log directory exists within ComfyUI's temp directory and returns its path."""
+    """
+    Ensures the API log directory exists within ComfyUI's temp directory
+    and returns its path.
+    """
     base_temp_dir = folder_paths.get_temp_directory()
     log_dir = os.path.join(base_temp_dir, "api_logs")
     try:
         os.makedirs(log_dir, exist_ok=True)
     except Exception as e:
-        logger.error("Error creating API log directory %s: %s", log_dir, str(e))
+        logger.error(f"Error creating API log directory {log_dir}: {e}")
         # Fallback to base temp directory if sub-directory creation fails
         return base_temp_dir
     return log_dir
 
-
-def _sanitize_filename_component(name: str) -> str:
-    if not name:
-        return "log"
-    sanitized = re.sub(r"[^A-Za-z0-9._-]+", "_", name)  # Replace disallowed characters with underscore
-    sanitized = sanitized.strip(" ._")  # Windows: trailing dots or spaces are not allowed
-    if not sanitized:
-        sanitized = "log"
-    return sanitized
-
-
-def _short_hash(*parts: str, length: int = 10) -> str:
-    return hashlib.sha1(("|".join(parts)).encode("utf-8")).hexdigest()[:length]
-
-
-def _build_log_filepath(log_dir: str, operation_id: str, request_url: str) -> str:
-    """Build log filepath. We keep it well under common path length limits aiming for <= 240 characters total."""
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
-    slug = _sanitize_filename_component(operation_id)  # Best-effort human-readable slug from operation_id
-    h = _short_hash(operation_id or "", request_url or "")  # Short hash ties log to the full operation and URL
-
-    # Compute how much room we have for the slug given the directory length
-    # Keep total path length reasonably below ~260 on Windows.
-    max_total_path = 240
-    prefix = f"{timestamp}_"
-    suffix = f"_{h}.log"
-    if not slug:
-        slug = "op"
-    max_filename_len = max(60, max_total_path - len(log_dir) - 1)
-    max_slug_len = max(8, max_filename_len - len(prefix) - len(suffix))
-    if len(slug) > max_slug_len:
-        slug = slug[:max_slug_len].rstrip(" ._-")
-    return os.path.join(log_dir, f"{prefix}{slug}{suffix}")
-
-
-def _format_data_for_logging(data: Any) -> str:
+def _format_data_for_logging(data):
     """Helper to format data (dict, str, bytes) for logging."""
     if isinstance(data, bytes):
         try:
-            return data.decode("utf-8")  # Try to decode as text
+            return data.decode('utf-8')  # Try to decode as text
         except UnicodeDecodeError:
             return f"[Binary data of length {len(data)} bytes]"
     elif isinstance(data, (dict, list)):
         try:
             return json.dumps(data, indent=2, ensure_ascii=False)
         except TypeError:
-            return str(data)  # Fallback for non-serializable objects
+            return str(data) # Fallback for non-serializable objects
     return str(data)
 
-
 def log_request_response(
     operation_id: str,
     request_method: str,
     request_url: str,
     request_headers: dict | None = None,
     request_params: dict | None = None,
-    request_data: Any = None,
+    request_data: any = None,
     response_status_code: int | None = None,
     response_headers: dict | None = None,
-    response_content: Any = None,
-    error_message: str | None = None,
+    response_content: any = None,
+    error_message: str | None = None
 ):
     """
     Logs API request and response details to a file in the temp/api_logs directory.
-    Filenames are sanitized and length-limited for cross-platform safety.
-    If we still fail to write, we fall back to appending into api.log.
     """
     log_dir = get_log_directory()
-    filepath = _build_log_filepath(log_dir, operation_id, request_url)
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+    filename = f"{timestamp}_{operation_id.replace('/', '_').replace(':', '_')}.log"
+    filepath = os.path.join(log_dir, filename)
+
+    log_content = []
 
-    log_content: list[str] = []
     log_content.append(f"Timestamp: {datetime.datetime.now().isoformat()}")
     log_content.append(f"Operation ID: {operation_id}")
     log_content.append("-" * 30 + " REQUEST " + "-" * 30)
@@ -106,7 +69,7 @@ def log_request_response(
         log_content.append(f"Headers:\n{_format_data_for_logging(request_headers)}")
     if request_params:
         log_content.append(f"Params:\n{_format_data_for_logging(request_params)}")
-    if request_data is not None:
+    if request_data:
         log_content.append(f"Data/Body:\n{_format_data_for_logging(request_data)}")
 
     log_content.append("\n" + "-" * 30 + " RESPONSE " + "-" * 30)
@@ -114,7 +77,7 @@ def log_request_response(
         log_content.append(f"Status Code: {response_status_code}")
     if response_headers:
         log_content.append(f"Headers:\n{_format_data_for_logging(response_headers)}")
-    if response_content is not None:
+    if response_content:
         log_content.append(f"Content:\n{_format_data_for_logging(response_content)}")
     if error_message:
         log_content.append(f"Error:\n{error_message}")
@@ -122,10 +85,9 @@ def log_request_response(
     try:
         with open(filepath, "w", encoding="utf-8") as f:
             f.write("\n".join(log_content))
-        logger.debug("API log saved to: %s", filepath)
+        logger.debug(f"API log saved to: {filepath}")
     except Exception as e:
-        logger.error("Error writing API log to %s: %s", filepath, str(e))
-
+        logger.error(f"Error writing API log to {filepath}: {e}")
 
 if __name__ == '__main__':
     # Example usage (for testing the logger directly)

comfy_api_nodes/apis/rodin_api.py

@@ -9,9 +9,8 @@ class Rodin3DGenerateRequest(BaseModel):
     seed: int = Field(..., description="seed_")
     tier: str = Field(..., description="Tier of generation.")
     material: str = Field(..., description="The material type.")
-    quality_override: int = Field(..., description="The poly count of the mesh.")
+    quality: str = Field(..., description="The generation quality of the mesh.")
     mesh_mode: str = Field(..., description="It controls the type of faces of generated models.")
-    TAPose: Optional[bool] = Field(None, description="")
 
 class GenerateJobsData(BaseModel):
     uuids: List[str] = Field(..., description="str LIST")
@@ -52,3 +51,7 @@ class RodinResourceItem(BaseModel):
 
 class Rodin3DDownloadResponse(BaseModel):
     list: List[RodinResourceItem] = Field(..., description="Source List")
+
+
+
+

comfy_api_nodes/nodes_bytedance.py

@@ -249,8 +249,8 @@ class ByteDanceImageNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=Text2ImageModelName,
-                    default=Text2ImageModelName.seedream_3,
+                    options=[model.value for model in Text2ImageModelName],
+                    default=Text2ImageModelName.seedream_3.value,
                     tooltip="Model name",
                 ),
                 comfy_io.String.Input(
@@ -382,8 +382,8 @@ class ByteDanceImageEditNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=Image2ImageModelName,
-                    default=Image2ImageModelName.seededit_3,
+                    options=[model.value for model in Image2ImageModelName],
+                    default=Image2ImageModelName.seededit_3.value,
                     tooltip="Model name",
                 ),
                 comfy_io.Image.Input(
@@ -676,8 +676,8 @@ class ByteDanceTextToVideoNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=Text2VideoModelName,
-                    default=Text2VideoModelName.seedance_1_pro,
+                    options=[model.value for model in Text2VideoModelName],
+                    default=Text2VideoModelName.seedance_1_pro.value,
                     tooltip="Model name",
                 ),
                 comfy_io.String.Input(
@@ -793,8 +793,8 @@ class ByteDanceImageToVideoNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=Image2VideoModelName,
-                    default=Image2VideoModelName.seedance_1_pro,
+                    options=[model.value for model in Image2VideoModelName],
+                    default=Image2VideoModelName.seedance_1_pro.value,
                     tooltip="Model name",
                 ),
                 comfy_io.String.Input(
@@ -920,7 +920,7 @@ class ByteDanceFirstLastFrameNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=[model.value for model in Image2VideoModelName],
+                    options=[Image2VideoModelName.seedance_1_lite.value],
                     default=Image2VideoModelName.seedance_1_lite.value,
                     tooltip="Model name",
                 ),

comfy_api_nodes/nodes_gemini.py

@@ -39,7 +39,6 @@ from comfy_api_nodes.apinode_utils import (
     tensor_to_base64_string,
     bytesio_to_image_tensor,
 )
-from comfy_api.util import VideoContainer, VideoCodec
 
 
 GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
@@ -311,7 +310,7 @@ class GeminiNode(ComfyNodeABC):
         Returns:
             List of GeminiPart objects containing the encoded video.
         """
-
+        from comfy_api.util import VideoContainer, VideoCodec
         base_64_string = video_to_base64_string(
             video_input,
             container_format=VideoContainer.MP4,
@@ -491,6 +490,7 @@ class GeminiInputFiles(ComfyNodeABC):
         # Use base64 string directly, not the data URI
         with open(file_path, "rb") as f:
             file_content = f.read()
+        import base64
         base64_str = base64.b64encode(file_content).decode("utf-8")
 
         return GeminiPart(

comfy_api_nodes/nodes_luma.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 from inspect import cleandoc
 from typing import Optional
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io as comfy_io
+from comfy.comfy_types.node_typing import IO, ComfyNodeABC
 from comfy_api.input_impl.video_types import VideoFromFile
 from comfy_api_nodes.apis.luma_api import (
     LumaImageModel,
@@ -52,186 +51,174 @@ def image_result_url_extractor(response: LumaGeneration):
 def video_result_url_extractor(response: LumaGeneration):
     return response.assets.video if hasattr(response, "assets") and hasattr(response.assets, "video") else None
 
-class LumaReferenceNode(comfy_io.ComfyNode):
+class LumaReferenceNode(ComfyNodeABC):
     """
     Holds an image and weight for use with Luma Generate Image node.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaReferenceNode",
-            display_name="Luma Reference",
-            category="api node/image/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input(
-                    "image",
-                    tooltip="Image to use as reference.",
-                ),
-                comfy_io.Float.Input(
-                    "weight",
-                    default=1.0,
-                    min=0.0,
-                    max=1.0,
-                    step=0.01,
-                    tooltip="Weight of image reference.",
-                ),
-                comfy_io.Custom(LumaIO.LUMA_REF).Input(
-                    "luma_ref",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Custom(LumaIO.LUMA_REF).Output(display_name="luma_ref")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-        )
+    RETURN_TYPES = (LumaIO.LUMA_REF,)
+    RETURN_NAMES = ("luma_ref",)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "create_luma_reference"
+    CATEGORY = "api node/image/Luma"
 
     @classmethod
-    def execute(
-        cls, image: torch.Tensor, weight: float, luma_ref: LumaReferenceChain = None
-    ) -> comfy_io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (
+                    IO.IMAGE,
+                    {
+                        "tooltip": "Image to use as reference.",
+                    },
+                ),
+                "weight": (
+                    IO.FLOAT,
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                        "tooltip": "Weight of image reference.",
+                    },
+                ),
+            },
+            "optional": {"luma_ref": (LumaIO.LUMA_REF,)},
+        }
+
+    def create_luma_reference(
+        self, image: torch.Tensor, weight: float, luma_ref: LumaReferenceChain = None
+    ):
         if luma_ref is not None:
             luma_ref = luma_ref.clone()
         else:
             luma_ref = LumaReferenceChain()
         luma_ref.add(LumaReference(image=image, weight=round(weight, 2)))
-        return comfy_io.NodeOutput(luma_ref)
+        return (luma_ref,)
 
 
-class LumaConceptsNode(comfy_io.ComfyNode):
+class LumaConceptsNode(ComfyNodeABC):
     """
     Holds one or more Camera Concepts for use with Luma Text to Video and Luma Image to Video nodes.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaConceptsNode",
-            display_name="Luma Concepts",
-            category="api node/video/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Combo.Input(
-                    "concept1",
-                    options=get_luma_concepts(include_none=True),
-                ),
-                comfy_io.Combo.Input(
-                    "concept2",
-                    options=get_luma_concepts(include_none=True),
-                ),
-                comfy_io.Combo.Input(
-                    "concept3",
-                    options=get_luma_concepts(include_none=True),
-                ),
-                comfy_io.Combo.Input(
-                    "concept4",
-                    options=get_luma_concepts(include_none=True),
-                ),
-                comfy_io.Custom(LumaIO.LUMA_CONCEPTS).Input(
-                    "luma_concepts",
-                    tooltip="Optional Camera Concepts to add to the ones chosen here.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Custom(LumaIO.LUMA_CONCEPTS).Output(display_name="luma_concepts")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-        )
+    RETURN_TYPES = (LumaIO.LUMA_CONCEPTS,)
+    RETURN_NAMES = ("luma_concepts",)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "create_concepts"
+    CATEGORY = "api node/video/Luma"
 
     @classmethod
-    def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "concept1": (get_luma_concepts(include_none=True),),
+                "concept2": (get_luma_concepts(include_none=True),),
+                "concept3": (get_luma_concepts(include_none=True),),
+                "concept4": (get_luma_concepts(include_none=True),),
+            },
+            "optional": {
+                "luma_concepts": (
+                    LumaIO.LUMA_CONCEPTS,
+                    {
+                        "tooltip": "Optional Camera Concepts to add to the ones chosen here."
+                    },
+                ),
+            },
+        }
+
+    def create_concepts(
+        self,
         concept1: str,
         concept2: str,
         concept3: str,
         concept4: str,
         luma_concepts: LumaConceptChain = None,
-    ) -> comfy_io.NodeOutput:
+    ):
         chain = LumaConceptChain(str_list=[concept1, concept2, concept3, concept4])
         if luma_concepts is not None:
             chain = luma_concepts.clone_and_merge(chain)
-        return comfy_io.NodeOutput(chain)
+        return (chain,)
 
 
-class LumaImageGenerationNode(comfy_io.ComfyNode):
+class LumaImageGenerationNode(ComfyNodeABC):
     """
     Generates images synchronously based on prompt and aspect ratio.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaImageNode",
-            display_name="Luma Text to Image",
-            category="api node/image/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the image generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=LumaImageModel,
-                ),
-                comfy_io.Combo.Input(
-                    "aspect_ratio",
-                    options=LumaAspectRatio,
-                    default=LumaAspectRatio.ratio_16_9,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    control_after_generate=True,
-                    tooltip="Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
-                ),
-                comfy_io.Float.Input(
-                    "style_image_weight",
-                    default=1.0,
-                    min=0.0,
-                    max=1.0,
-                    step=0.01,
-                    tooltip="Weight of style image. Ignored if no style_image provided.",
-                ),
-                comfy_io.Custom(LumaIO.LUMA_REF).Input(
-                    "image_luma_ref",
-                    tooltip="Luma Reference node connection to influence generation with input images; up to 4 images can be considered.",
-                    optional=True,
-                ),
-                comfy_io.Image.Input(
-                    "style_image",
-                    tooltip="Style reference image; only 1 image will be used.",
-                    optional=True,
-                ),
-                comfy_io.Image.Input(
-                    "character_image",
-                    tooltip="Character reference images; can be a batch of multiple, up to 4 images can be considered.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Image.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.IMAGE,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/image/Luma"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the image generation",
+                    },
+                ),
+                "model": ([model.value for model in LumaImageModel],),
+                "aspect_ratio": (
+                    [ratio.value for ratio in LumaAspectRatio],
+                    {
+                        "default": LumaAspectRatio.ratio_16_9,
+                    },
+                ),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                    },
+                ),
+                "style_image_weight": (
+                    IO.FLOAT,
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                        "tooltip": "Weight of style image. Ignored if no style_image provided.",
+                    },
+                ),
+            },
+            "optional": {
+                "image_luma_ref": (
+                    LumaIO.LUMA_REF,
+                    {
+                        "tooltip": "Luma Reference node connection to influence generation with input images; up to 4 images can be considered."
+                    },
+                ),
+                "style_image": (
+                    IO.IMAGE,
+                    {"tooltip": "Style reference image; only 1 image will be used."},
+                ),
+                "character_image": (
+                    IO.IMAGE,
+                    {
+                        "tooltip": "Character reference images; can be a batch of multiple, up to 4 images can be considered."
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         prompt: str,
         model: str,
         aspect_ratio: str,
@@ -240,29 +227,27 @@ class LumaImageGenerationNode(comfy_io.ComfyNode):
         image_luma_ref: LumaReferenceChain = None,
         style_image: torch.Tensor = None,
         character_image: torch.Tensor = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: str = None,
+        **kwargs,
+    ):
         validate_string(prompt, strip_whitespace=True, min_length=3)
-        auth_kwargs = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
         # handle image_luma_ref
         api_image_ref = None
         if image_luma_ref is not None:
-            api_image_ref = await cls._convert_luma_refs(
-                image_luma_ref, max_refs=4, auth_kwargs=auth_kwargs,
+            api_image_ref = await self._convert_luma_refs(
+                image_luma_ref, max_refs=4, auth_kwargs=kwargs,
             )
         # handle style_luma_ref
         api_style_ref = None
         if style_image is not None:
-            api_style_ref = await cls._convert_style_image(
-                style_image, weight=style_image_weight, auth_kwargs=auth_kwargs,
+            api_style_ref = await self._convert_style_image(
+                style_image, weight=style_image_weight, auth_kwargs=kwargs,
             )
         # handle character_ref images
         character_ref = None
         if character_image is not None:
             download_urls = await upload_images_to_comfyapi(
-                character_image, max_images=4, auth_kwargs=auth_kwargs,
+                character_image, max_images=4, auth_kwargs=kwargs,
             )
             character_ref = LumaCharacterRef(
                 identity0=LumaImageIdentity(images=download_urls)
@@ -283,7 +268,7 @@ class LumaImageGenerationNode(comfy_io.ComfyNode):
                 style_ref=api_style_ref,
                 character_ref=character_ref,
             ),
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_api: LumaGeneration = await operation.execute()
 
@@ -298,19 +283,18 @@ class LumaImageGenerationNode(comfy_io.ComfyNode):
             failed_statuses=[LumaState.failed],
             status_extractor=lambda x: x.state,
             result_url_extractor=image_result_url_extractor,
-            node_id=cls.hidden.unique_id,
-            auth_kwargs=auth_kwargs,
+            node_id=unique_id,
+            auth_kwargs=kwargs,
         )
         response_poll = await operation.execute()
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.assets.image) as img_response:
                 img = process_image_response(await img_response.content.read())
-        return comfy_io.NodeOutput(img)
+        return (img,)
 
-    @classmethod
     async def _convert_luma_refs(
-        cls, luma_ref: LumaReferenceChain, max_refs: int, auth_kwargs: Optional[dict[str,str]] = None
+        self, luma_ref: LumaReferenceChain, max_refs: int, auth_kwargs: Optional[dict[str,str]] = None
     ):
         luma_urls = []
         ref_count = 0
@@ -324,84 +308,82 @@ class LumaImageGenerationNode(comfy_io.ComfyNode):
                 break
         return luma_ref.create_api_model(download_urls=luma_urls, max_refs=max_refs)
 
-    @classmethod
     async def _convert_style_image(
-        cls, style_image: torch.Tensor, weight: float, auth_kwargs: Optional[dict[str,str]] = None
+        self, style_image: torch.Tensor, weight: float, auth_kwargs: Optional[dict[str,str]] = None
     ):
         chain = LumaReferenceChain(
             first_ref=LumaReference(image=style_image, weight=weight)
         )
-        return await cls._convert_luma_refs(chain, max_refs=1, auth_kwargs=auth_kwargs)
+        return await self._convert_luma_refs(chain, max_refs=1, auth_kwargs=auth_kwargs)
 
 
-class LumaImageModifyNode(comfy_io.ComfyNode):
+class LumaImageModifyNode(ComfyNodeABC):
     """
     Modifies images synchronously based on prompt and aspect ratio.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaImageModifyNode",
-            display_name="Luma Image to Image",
-            category="api node/image/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input(
-                    "image",
-                ),
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the image generation",
-                ),
-                comfy_io.Float.Input(
-                    "image_weight",
-                    default=0.1,
-                    min=0.0,
-                    max=0.98,
-                    step=0.01,
-                    tooltip="Weight of the image; the closer to 1.0, the less the image will be modified.",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=LumaImageModel,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    control_after_generate=True,
-                    tooltip="Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
-                ),
-            ],
-            outputs=[comfy_io.Image.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.IMAGE,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/image/Luma"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (IO.IMAGE,),
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the image generation",
+                    },
+                ),
+                "image_weight": (
+                    IO.FLOAT,
+                    {
+                        "default": 0.1,
+                        "min": 0.0,
+                        "max": 0.98,
+                        "step": 0.01,
+                        "tooltip": "Weight of the image; the closer to 1.0, the less the image will be modified.",
+                    },
+                ),
+                "model": ([model.value for model in LumaImageModel],),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                    },
+                ),
+            },
+            "optional": {},
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         prompt: str,
         model: str,
         image: torch.Tensor,
         image_weight: float,
         seed,
-    ) -> comfy_io.NodeOutput:
-        auth_kwargs = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
+        unique_id: str = None,
+        **kwargs,
+    ):
         # first, upload image
         download_urls = await upload_images_to_comfyapi(
-            image, max_images=1, auth_kwargs=auth_kwargs,
+            image, max_images=1, auth_kwargs=kwargs,
         )
         image_url = download_urls[0]
         # next, make Luma call with download url provided
@@ -419,7 +401,7 @@ class LumaImageModifyNode(comfy_io.ComfyNode):
                     url=image_url, weight=round(max(min(1.0-image_weight, 0.98), 0.0), 2)
                 ),
             ),
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_api: LumaGeneration = await operation.execute()
 
@@ -434,84 +416,88 @@ class LumaImageModifyNode(comfy_io.ComfyNode):
             failed_statuses=[LumaState.failed],
             status_extractor=lambda x: x.state,
             result_url_extractor=image_result_url_extractor,
-            node_id=cls.hidden.unique_id,
-            auth_kwargs=auth_kwargs,
+            node_id=unique_id,
+            auth_kwargs=kwargs,
         )
         response_poll = await operation.execute()
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.assets.image) as img_response:
                 img = process_image_response(await img_response.content.read())
-        return comfy_io.NodeOutput(img)
+        return (img,)
 
 
-class LumaTextToVideoGenerationNode(comfy_io.ComfyNode):
+class LumaTextToVideoGenerationNode(ComfyNodeABC):
     """
     Generates videos synchronously based on prompt and output_size.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaVideoNode",
-            display_name="Luma Text to Video",
-            category="api node/video/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=LumaVideoModel,
-                ),
-                comfy_io.Combo.Input(
-                    "aspect_ratio",
-                    options=LumaAspectRatio,
-                    default=LumaAspectRatio.ratio_16_9,
-                ),
-                comfy_io.Combo.Input(
-                    "resolution",
-                    options=LumaVideoOutputResolution,
-                    default=LumaVideoOutputResolution.res_540p,
-                ),
-                comfy_io.Combo.Input(
-                    "duration",
-                    options=LumaVideoModelOutputDuration,
-                ),
-                comfy_io.Boolean.Input(
-                    "loop",
-                    default=False,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    control_after_generate=True,
-                    tooltip="Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
-                ),
-                comfy_io.Custom(LumaIO.LUMA_CONCEPTS).Input(
-                    "luma_concepts",
-                    tooltip="Optional Camera Concepts to dictate camera motion via the Luma Concepts node.",
-                    optional=True,
-                )
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.VIDEO,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/video/Luma"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the video generation",
+                    },
+                ),
+                "model": ([model.value for model in LumaVideoModel],),
+                "aspect_ratio": (
+                    [ratio.value for ratio in LumaAspectRatio],
+                    {
+                        "default": LumaAspectRatio.ratio_16_9,
+                    },
+                ),
+                "resolution": (
+                    [resolution.value for resolution in LumaVideoOutputResolution],
+                    {
+                        "default": LumaVideoOutputResolution.res_540p,
+                    },
+                ),
+                "duration": ([dur.value for dur in LumaVideoModelOutputDuration],),
+                "loop": (
+                    IO.BOOLEAN,
+                    {
+                        "default": False,
+                    },
+                ),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                    },
+                ),
+            },
+            "optional": {
+                "luma_concepts": (
+                    LumaIO.LUMA_CONCEPTS,
+                    {
+                        "tooltip": "Optional Camera Concepts to dictate camera motion via the Luma Concepts node."
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         prompt: str,
         model: str,
         aspect_ratio: str,
@@ -520,15 +506,13 @@ class LumaTextToVideoGenerationNode(comfy_io.ComfyNode):
         loop: bool,
         seed,
         luma_concepts: LumaConceptChain = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: str = None,
+        **kwargs,
+    ):
         validate_string(prompt, strip_whitespace=False, min_length=3)
         duration = duration if model != LumaVideoModel.ray_1_6 else None
         resolution = resolution if model != LumaVideoModel.ray_1_6 else None
 
-        auth_kwargs = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/luma/generations",
@@ -545,12 +529,12 @@ class LumaTextToVideoGenerationNode(comfy_io.ComfyNode):
                 loop=loop,
                 concepts=luma_concepts.create_api_model() if luma_concepts else None,
             ),
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_api: LumaGeneration = await operation.execute()
 
-        if cls.hidden.unique_id:
-            PromptServer.instance.send_progress_text(f"Luma video generation started: {response_api.id}", cls.hidden.unique_id)
+        if unique_id:
+            PromptServer.instance.send_progress_text(f"Luma video generation started: {response_api.id}", unique_id)
 
         operation = PollingOperation(
             poll_endpoint=ApiEndpoint(
@@ -563,94 +547,90 @@ class LumaTextToVideoGenerationNode(comfy_io.ComfyNode):
             failed_statuses=[LumaState.failed],
             status_extractor=lambda x: x.state,
             result_url_extractor=video_result_url_extractor,
-            node_id=cls.hidden.unique_id,
+            node_id=unique_id,
             estimated_duration=LUMA_T2V_AVERAGE_DURATION,
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_poll = await operation.execute()
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.assets.video) as vid_response:
-                return comfy_io.NodeOutput(VideoFromFile(BytesIO(await vid_response.content.read())))
+                return (VideoFromFile(BytesIO(await vid_response.content.read())),)
 
 
-class LumaImageToVideoGenerationNode(comfy_io.ComfyNode):
+class LumaImageToVideoGenerationNode(ComfyNodeABC):
     """
     Generates videos synchronously based on prompt, input images, and output_size.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="LumaImageToVideoNode",
-            display_name="Luma Image to Video",
-            category="api node/video/Luma",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=LumaVideoModel,
-                ),
-                # comfy_io.Combo.Input(
-                #     "aspect_ratio",
-                #     options=[ratio.value for ratio in LumaAspectRatio],
-                #     default=LumaAspectRatio.ratio_16_9,
-                # ),
-                comfy_io.Combo.Input(
-                    "resolution",
-                    options=LumaVideoOutputResolution,
-                    default=LumaVideoOutputResolution.res_540p,
-                ),
-                comfy_io.Combo.Input(
-                    "duration",
-                    options=[dur.value for dur in LumaVideoModelOutputDuration],
-                ),
-                comfy_io.Boolean.Input(
-                    "loop",
-                    default=False,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    control_after_generate=True,
-                    tooltip="Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
-                ),
-                comfy_io.Image.Input(
-                    "first_image",
-                    tooltip="First frame of generated video.",
-                    optional=True,
-                ),
-                comfy_io.Image.Input(
-                    "last_image",
-                    tooltip="Last frame of generated video.",
-                    optional=True,
-                ),
-                comfy_io.Custom(LumaIO.LUMA_CONCEPTS).Input(
-                    "luma_concepts",
-                    tooltip="Optional Camera Concepts to dictate camera motion via the Luma Concepts node.",
-                    optional=True,
-                )
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.VIDEO,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/video/Luma"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the video generation",
+                    },
+                ),
+                "model": ([model.value for model in LumaVideoModel],),
+                # "aspect_ratio": ([ratio.value for ratio in LumaAspectRatio], {
+                #     "default": LumaAspectRatio.ratio_16_9,
+                # }),
+                "resolution": (
+                    [resolution.value for resolution in LumaVideoOutputResolution],
+                    {
+                        "default": LumaVideoOutputResolution.res_540p,
+                    },
+                ),
+                "duration": ([dur.value for dur in LumaVideoModelOutputDuration],),
+                "loop": (
+                    IO.BOOLEAN,
+                    {
+                        "default": False,
+                    },
+                ),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                    },
+                ),
+            },
+            "optional": {
+                "first_image": (
+                    IO.IMAGE,
+                    {"tooltip": "First frame of generated video."},
+                ),
+                "last_image": (IO.IMAGE, {"tooltip": "Last frame of generated video."}),
+                "luma_concepts": (
+                    LumaIO.LUMA_CONCEPTS,
+                    {
+                        "tooltip": "Optional Camera Concepts to dictate camera motion via the Luma Concepts node."
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         prompt: str,
         model: str,
         resolution: str,
@@ -660,16 +640,14 @@ class LumaImageToVideoGenerationNode(comfy_io.ComfyNode):
         first_image: torch.Tensor = None,
         last_image: torch.Tensor = None,
         luma_concepts: LumaConceptChain = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: str = None,
+        **kwargs,
+    ):
         if first_image is None and last_image is None:
             raise Exception(
                 "At least one of first_image and last_image requires an input."
             )
-        auth_kwargs = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        keyframes = await cls._convert_to_keyframes(first_image, last_image, auth_kwargs=auth_kwargs)
+        keyframes = await self._convert_to_keyframes(first_image, last_image, auth_kwargs=kwargs)
         duration = duration if model != LumaVideoModel.ray_1_6 else None
         resolution = resolution if model != LumaVideoModel.ray_1_6 else None
 
@@ -690,12 +668,12 @@ class LumaImageToVideoGenerationNode(comfy_io.ComfyNode):
                 keyframes=keyframes,
                 concepts=luma_concepts.create_api_model() if luma_concepts else None,
             ),
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_api: LumaGeneration = await operation.execute()
 
-        if cls.hidden.unique_id:
-            PromptServer.instance.send_progress_text(f"Luma video generation started: {response_api.id}", cls.hidden.unique_id)
+        if unique_id:
+            PromptServer.instance.send_progress_text(f"Luma video generation started: {response_api.id}", unique_id)
 
         operation = PollingOperation(
             poll_endpoint=ApiEndpoint(
@@ -708,19 +686,18 @@ class LumaImageToVideoGenerationNode(comfy_io.ComfyNode):
             failed_statuses=[LumaState.failed],
             status_extractor=lambda x: x.state,
             result_url_extractor=video_result_url_extractor,
-            node_id=cls.hidden.unique_id,
+            node_id=unique_id,
             estimated_duration=LUMA_I2V_AVERAGE_DURATION,
-            auth_kwargs=auth_kwargs,
+            auth_kwargs=kwargs,
         )
         response_poll = await operation.execute()
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.assets.video) as vid_response:
-                return comfy_io.NodeOutput(VideoFromFile(BytesIO(await vid_response.content.read())))
+                return (VideoFromFile(BytesIO(await vid_response.content.read())),)
 
-    @classmethod
     async def _convert_to_keyframes(
-        cls,
+        self,
         first_image: torch.Tensor = None,
         last_image: torch.Tensor = None,
         auth_kwargs: Optional[dict[str,str]] = None,
@@ -742,18 +719,23 @@ class LumaImageToVideoGenerationNode(comfy_io.ComfyNode):
         return LumaKeyframes(frame0=frame0, frame1=frame1)
 
 
-class LumaExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            LumaImageGenerationNode,
-            LumaImageModifyNode,
-            LumaTextToVideoGenerationNode,
-            LumaImageToVideoGenerationNode,
-            LumaReferenceNode,
-            LumaConceptsNode,
-        ]
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "LumaImageNode": LumaImageGenerationNode,
+    "LumaImageModifyNode": LumaImageModifyNode,
+    "LumaVideoNode": LumaTextToVideoGenerationNode,
+    "LumaImageToVideoNode": LumaImageToVideoGenerationNode,
+    "LumaReferenceNode": LumaReferenceNode,
+    "LumaConceptsNode": LumaConceptsNode,
+}
 
-
-async def comfy_entrypoint() -> LumaExtension:
-    return LumaExtension()
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LumaImageNode": "Luma Text to Image",
+    "LumaImageModifyNode": "Luma Image to Image",
+    "LumaVideoNode": "Luma Text to Video",
+    "LumaImageToVideoNode": "Luma Image to Video",
+    "LumaReferenceNode": "Luma Reference",
+    "LumaConceptsNode": "Luma Concepts",
+}

comfy_api_nodes/nodes_minimax.py

@@ -500,7 +500,7 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
             raise Exception(
                 f"No video was found in the response. Full response: {file_result.model_dump()}"
             )
-        logging.info("Generated video URL: %s", file_url)
+        logging.info(f"Generated video URL: {file_url}")
         if cls.hidden.unique_id:
             if hasattr(file_result.file, "backup_download_url"):
                 message = f"Result URL: {file_url}\nBackup URL: {file_result.file.backup_download_url}"

comfy_api_nodes/nodes_moonvalley.py

@@ -2,7 +2,11 @@ import logging
 from typing import Any, Callable, Optional, TypeVar
 import torch
 from typing_extensions import override
-from comfy_api_nodes.util.validation_utils import validate_image_dimensions
+from comfy_api_nodes.util.validation_utils import (
+    get_image_dimensions,
+    validate_image_dimensions,
+)
+
 
 from comfy_api_nodes.apis import (
     MoonvalleyTextToVideoRequest,
@@ -128,6 +132,47 @@ def validate_prompts(
     return True
 
 
+def validate_input_media(width, height, with_frame_conditioning, num_frames_in=None):
+    # inference validation
+    # T = num_frames
+    # in all cases, the following must be true: T divisible by 16 and H,W by 8. in addition...
+    # with image conditioning: H*W must be divisible by 8192
+    # without image conditioning: T divisible by 32
+    if num_frames_in and not num_frames_in % 16 == 0:
+        return False, ("The input video total frame count must be divisible by 16!")
+
+    if height % 8 != 0 or width % 8 != 0:
+        return False, (
+            f"Height ({height}) and width ({width}) must be " "divisible by 8"
+        )
+
+    if with_frame_conditioning:
+        if (height * width) % 8192 != 0:
+            return False, (
+                f"Height * width ({height * width}) must be "
+                "divisible by 8192 for frame conditioning"
+            )
+    else:
+        if num_frames_in and not num_frames_in % 32 == 0:
+            return False, ("The input video total frame count must be divisible by 32!")
+
+
+def validate_input_image(
+    image: torch.Tensor, with_frame_conditioning: bool = False
+) -> None:
+    """
+    Validates the input image adheres to the expectations of the API:
+    - The image resolution should not be less than 300*300px
+    - The aspect ratio of the image should be between 1:2.5 ~ 2.5:1
+
+    """
+    height, width = get_image_dimensions(image)
+    validate_input_media(width, height, with_frame_conditioning)
+    validate_image_dimensions(
+        image, min_width=300, min_height=300, max_height=MAX_HEIGHT, max_width=MAX_WIDTH
+    )
+
+
 def validate_video_to_video_input(video: VideoInput) -> VideoInput:
     """
     Validates and processes video input for Moonvalley Video-to-Video generation.
@@ -237,7 +282,7 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
         audio_stream = None
 
         for stream in input_container.streams:
-            logging.info("Found stream: type=%s, class=%s", stream.type, type(stream))
+            logging.info(f"Found stream: type={stream.type}, class={type(stream)}")
             if isinstance(stream, av.VideoStream):
                 # Create output video stream with same parameters
                 video_stream = output_container.add_stream(
@@ -247,7 +292,7 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
                 video_stream.height = stream.height
                 video_stream.pix_fmt = "yuv420p"
                 logging.info(
-                    "Added video stream: %sx%s @ %sfps", stream.width, stream.height, stream.average_rate
+                    f"Added video stream: {stream.width}x{stream.height} @ {stream.average_rate}fps"
                 )
             elif isinstance(stream, av.AudioStream):
                 # Create output audio stream with same parameters
@@ -256,7 +301,9 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
                 )
                 audio_stream.sample_rate = stream.sample_rate
                 audio_stream.layout = stream.layout
-                logging.info("Added audio stream: %sHz, %s channels", stream.sample_rate, stream.channels)
+                logging.info(
+                    f"Added audio stream: {stream.sample_rate}Hz, {stream.channels} channels"
+                )
 
         # Calculate target frame count that's divisible by 16
         fps = input_container.streams.video[0].average_rate
@@ -286,7 +333,9 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
             for packet in video_stream.encode():
                 output_container.mux(packet)
 
-            logging.info("Encoded %s video frames (target: %s)", frame_count, target_frames)
+            logging.info(
+                f"Encoded {frame_count} video frames (target: {target_frames})"
+            )
 
         # Decode and re-encode audio frames
         if audio_stream:
@@ -304,7 +353,7 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
             for packet in audio_stream.encode():
                 output_container.mux(packet)
 
-            logging.info("Encoded %s audio frames", audio_frame_count)
+            logging.info(f"Encoded {audio_frame_count} audio frames")
 
         # Close containers
         output_container.close()
@@ -331,7 +380,7 @@ def parse_width_height_from_res(resolution: str):
         "1:1 (1152 x 1152)": {"width": 1152, "height": 1152},
         "4:3 (1536 x 1152)": {"width": 1536, "height": 1152},
         "3:4 (1152 x 1536)": {"width": 1152, "height": 1536},
-        # "21:9 (2560 x 1080)": {"width": 2560, "height": 1080},
+        "21:9 (2560 x 1080)": {"width": 2560, "height": 1080},
     }
     return res_map.get(resolution, {"width": 1920, "height": 1080})
 
@@ -384,11 +433,11 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
                     "negative_prompt",
                     multiline=True,
                     default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                    "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                    "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                    "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                    "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                    "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
+                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
+                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
+                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
+                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
+                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
                     tooltip="Negative prompt text",
                 ),
                 comfy_io.Combo.Input(
@@ -399,14 +448,14 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
                         "1:1 (1152 x 1152)",
                         "4:3 (1536 x 1152)",
                         "3:4 (1152 x 1536)",
-                        # "21:9 (2560 x 1080)",
+                        "21:9 (2560 x 1080)",
                     ],
                     default="16:9 (1920 x 1080)",
                     tooltip="Resolution of the output video",
                 ),
                 comfy_io.Float.Input(
                     "prompt_adherence",
-                    default=4.5,
+                    default=10.0,
                     min=1.0,
                     max=20.0,
                     step=1.0,
@@ -420,11 +469,10 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
                     step=1,
                     display_mode=comfy_io.NumberDisplay.number,
                     tooltip="Random seed value",
-                    control_after_generate=True,
                 ),
                 comfy_io.Int.Input(
                     "steps",
-                    default=33,
+                    default=100,
                     min=1,
                     max=100,
                     step=1,
@@ -451,7 +499,7 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
         seed: int,
         steps: int,
     ) -> comfy_io.NodeOutput:
-        validate_image_dimensions(image, min_width=300, min_height=300, max_height=MAX_HEIGHT, max_width=MAX_WIDTH)
+        validate_input_image(image, True)
         validate_prompts(prompt, negative_prompt, MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
         width_height = parse_width_height_from_res(resolution)
 
@@ -465,11 +513,12 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
             steps=steps,
             seed=seed,
             guidance_scale=prompt_adherence,
+            num_frames=128,
             width=width_height["width"],
             height=width_height["height"],
             use_negative_prompts=True,
         )
-
+        """Upload image to comfy backend to have a URL available for further processing"""
         # Get MIME type from tensor - assuming PNG format for image tensors
         mime_type = "image/png"
 
@@ -522,11 +571,11 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
                     "negative_prompt",
                     multiline=True,
                     default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                    "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                    "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                    "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                    "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                    "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
+                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
+                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
+                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
+                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
+                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
                     tooltip="Negative prompt text",
                 ),
                 comfy_io.Int.Input(
@@ -542,7 +591,7 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
                 comfy_io.Video.Input(
                     "video",
                     tooltip="The reference video used to generate the output video. Must be at least 5 seconds long. "
-                    "Videos longer than 5s will be automatically trimmed. Only MP4 format supported.",
+                            "Videos longer than 5s will be automatically trimmed. Only MP4 format supported.",
                 ),
                 comfy_io.Combo.Input(
                     "control_type",
@@ -559,15 +608,6 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
                     tooltip="Only used if control_type is 'Motion Transfer'",
                     optional=True,
                 ),
-                comfy_io.Int.Input(
-                    "steps",
-                    default=33,
-                    min=1,
-                    max=100,
-                    step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    tooltip="Number of inference steps",
-                ),
             ],
             outputs=[comfy_io.Video.Output()],
             hidden=[
@@ -587,8 +627,6 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
         video: Optional[VideoInput] = None,
         control_type: str = "Motion Transfer",
         motion_intensity: Optional[int] = 100,
-        steps=33,
-        prompt_adherence=4.5,
     ) -> comfy_io.NodeOutput:
         auth = {
             "auth_token": cls.hidden.auth_token_comfy_org,
@@ -598,6 +636,7 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
         validated_video = validate_video_to_video_input(video)
         video_url = await upload_video_to_comfyapi(validated_video, auth_kwargs=auth)
 
+        """Validate prompts and inference input"""
         validate_prompts(prompt, negative_prompt)
 
         # Only include motion_intensity for Motion Transfer
@@ -609,8 +648,6 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
             negative_prompt=negative_prompt,
             seed=seed,
             control_params=control_params,
-            steps=steps,
-            guidance_scale=prompt_adherence,
         )
 
         control = parse_control_parameter(control_type)
@@ -662,11 +699,11 @@ class MoonvalleyTxt2VideoNode(comfy_io.ComfyNode):
                     "negative_prompt",
                     multiline=True,
                     default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                    "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                    "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                    "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                    "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                    "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
+                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
+                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
+                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
+                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
+                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
                     tooltip="Negative prompt text",
                 ),
                 comfy_io.Combo.Input(
@@ -684,7 +721,7 @@ class MoonvalleyTxt2VideoNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Float.Input(
                     "prompt_adherence",
-                    default=4.0,
+                    default=10.0,
                     min=1.0,
                     max=20.0,
                     step=1.0,
@@ -697,12 +734,11 @@ class MoonvalleyTxt2VideoNode(comfy_io.ComfyNode):
                     max=4294967295,
                     step=1,
                     display_mode=comfy_io.NumberDisplay.number,
-                    control_after_generate=True,
                     tooltip="Random seed value",
                 ),
                 comfy_io.Int.Input(
                     "steps",
-                    default=33,
+                    default=100,
                     min=1,
                     max=100,
                     step=1,

comfy_api_nodes/nodes_pixverse.py

@@ -1,7 +1,5 @@
 from inspect import cleandoc
 from typing import Optional
-from typing_extensions import override
-from io import BytesIO
 from comfy_api_nodes.apis.pixverse_api import (
     PixverseTextVideoRequest,
     PixverseImageVideoRequest,
@@ -28,11 +26,12 @@ from comfy_api_nodes.apinode_utils import (
     tensor_to_bytesio,
     validate_string,
 )
+from comfy.comfy_types.node_typing import IO, ComfyNodeABC
 from comfy_api.input_impl import VideoFromFile
-from comfy_api.latest import ComfyExtension, io as comfy_io
 
 import torch
 import aiohttp
+from io import BytesIO
 
 
 AVERAGE_DURATION_T2V = 32
@@ -73,101 +72,100 @@ async def upload_image_to_pixverse(image: torch.Tensor, auth_kwargs=None):
     return response_upload.Resp.img_id
 
 
-class PixverseTemplateNode(comfy_io.ComfyNode):
+class PixverseTemplateNode:
     """
     Select template for PixVerse Video generation.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="PixverseTemplateNode",
-            display_name="PixVerse Template",
-            category="api node/video/PixVerse",
-            inputs=[
-                comfy_io.Combo.Input("template", options=list(pixverse_templates.keys())),
-            ],
-            outputs=[comfy_io.Custom(PixverseIO.TEMPLATE).Output(display_name="pixverse_template")],
-        )
+    RETURN_TYPES = (PixverseIO.TEMPLATE,)
+    RETURN_NAMES = ("pixverse_template",)
+    FUNCTION = "create_template"
+    CATEGORY = "api node/video/PixVerse"
 
     @classmethod
-    def execute(cls, template: str) -> comfy_io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "template": (list(pixverse_templates.keys()),),
+            }
+        }
+
+    def create_template(self, template: str):
         template_id = pixverse_templates.get(template, None)
         if template_id is None:
             raise Exception(f"Template '{template}' is not recognized.")
         # just return the integer
-        return comfy_io.NodeOutput(template_id)
+        return (template_id,)
 
 
-class PixverseTextToVideoNode(comfy_io.ComfyNode):
+class PixverseTextToVideoNode(ComfyNodeABC):
     """
     Generates videos based on prompt and output_size.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="PixverseTextToVideoNode",
-            display_name="PixVerse Text to Video",
-            category="api node/video/PixVerse",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "aspect_ratio",
-                    options=PixverseAspectRatio,
-                ),
-                comfy_io.Combo.Input(
-                    "quality",
-                    options=PixverseQuality,
-                    default=PixverseQuality.res_540p,
-                ),
-                comfy_io.Combo.Input(
-                    "duration_seconds",
-                    options=PixverseDuration,
-                ),
-                comfy_io.Combo.Input(
-                    "motion_mode",
-                    options=PixverseMotionMode,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=2147483647,
-                    control_after_generate=True,
-                    tooltip="Seed for video generation.",
-                ),
-                comfy_io.String.Input(
-                    "negative_prompt",
-                    default="",
-                    multiline=True,
-                    tooltip="An optional text description of undesired elements on an image.",
-                    optional=True,
-                ),
-                comfy_io.Custom(PixverseIO.TEMPLATE).Input(
-                    "pixverse_template",
-                    tooltip="An optional template to influence style of generation, created by the PixVerse Template node.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.VIDEO,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/video/PixVerse"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the video generation",
+                    },
+                ),
+                "aspect_ratio": ([ratio.value for ratio in PixverseAspectRatio],),
+                "quality": (
+                    [resolution.value for resolution in PixverseQuality],
+                    {
+                        "default": PixverseQuality.res_540p,
+                    },
+                ),
+                "duration_seconds": ([dur.value for dur in PixverseDuration],),
+                "motion_mode": ([mode.value for mode in PixverseMotionMode],),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2147483647,
+                        "control_after_generate": True,
+                        "tooltip": "Seed for video generation.",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    IO.STRING,
+                    {
+                        "default": "",
+                        "forceInput": True,
+                        "tooltip": "An optional text description of undesired elements on an image.",
+                    },
+                ),
+                "pixverse_template": (
+                    PixverseIO.TEMPLATE,
+                    {
+                        "tooltip": "An optional template to influence style of generation, created by the PixVerse Template node."
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         prompt: str,
         aspect_ratio: str,
         quality: str,
@@ -176,7 +174,9 @@ class PixverseTextToVideoNode(comfy_io.ComfyNode):
         seed,
         negative_prompt: str = None,
         pixverse_template: int = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: Optional[str] = None,
+        **kwargs,
+    ):
         validate_string(prompt, strip_whitespace=False)
         # 1080p is limited to 5 seconds duration
         # only normal motion_mode supported for 1080p or for non-5 second duration
@@ -186,10 +186,6 @@ class PixverseTextToVideoNode(comfy_io.ComfyNode):
         elif duration_seconds != PixverseDuration.dur_5:
             motion_mode = PixverseMotionMode.normal
 
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/pixverse/video/text/generate",
@@ -207,7 +203,7 @@ class PixverseTextToVideoNode(comfy_io.ComfyNode):
                 template_id=pixverse_template,
                 seed=seed,
             ),
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         response_api = await operation.execute()
 
@@ -228,8 +224,8 @@ class PixverseTextToVideoNode(comfy_io.ComfyNode):
                 PixverseStatus.deleted,
             ],
             status_extractor=lambda x: x.Resp.status,
-            auth_kwargs=auth,
-            node_id=cls.hidden.unique_id,
+            auth_kwargs=kwargs,
+            node_id=unique_id,
             result_url_extractor=get_video_url_from_response,
             estimated_duration=AVERAGE_DURATION_T2V,
         )
@@ -237,75 +233,77 @@ class PixverseTextToVideoNode(comfy_io.ComfyNode):
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.Resp.url) as vid_response:
-                return comfy_io.NodeOutput(VideoFromFile(BytesIO(await vid_response.content.read())))
+                return (VideoFromFile(BytesIO(await vid_response.content.read())),)
 
 
-class PixverseImageToVideoNode(comfy_io.ComfyNode):
+class PixverseImageToVideoNode(ComfyNodeABC):
     """
     Generates videos based on prompt and output_size.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="PixverseImageToVideoNode",
-            display_name="PixVerse Image to Video",
-            category="api node/video/PixVerse",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("image"),
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "quality",
-                    options=PixverseQuality,
-                    default=PixverseQuality.res_540p,
-                ),
-                comfy_io.Combo.Input(
-                    "duration_seconds",
-                    options=PixverseDuration,
-                ),
-                comfy_io.Combo.Input(
-                    "motion_mode",
-                    options=PixverseMotionMode,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=2147483647,
-                    control_after_generate=True,
-                    tooltip="Seed for video generation.",
-                ),
-                comfy_io.String.Input(
-                    "negative_prompt",
-                    default="",
-                    multiline=True,
-                    tooltip="An optional text description of undesired elements on an image.",
-                    optional=True,
-                ),
-                comfy_io.Custom(PixverseIO.TEMPLATE).Input(
-                    "pixverse_template",
-                    tooltip="An optional template to influence style of generation, created by the PixVerse Template node.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.VIDEO,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/video/PixVerse"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (IO.IMAGE,),
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the video generation",
+                    },
+                ),
+                "quality": (
+                    [resolution.value for resolution in PixverseQuality],
+                    {
+                        "default": PixverseQuality.res_540p,
+                    },
+                ),
+                "duration_seconds": ([dur.value for dur in PixverseDuration],),
+                "motion_mode": ([mode.value for mode in PixverseMotionMode],),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2147483647,
+                        "control_after_generate": True,
+                        "tooltip": "Seed for video generation.",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    IO.STRING,
+                    {
+                        "default": "",
+                        "forceInput": True,
+                        "tooltip": "An optional text description of undesired elements on an image.",
+                    },
+                ),
+                "pixverse_template": (
+                    PixverseIO.TEMPLATE,
+                    {
+                        "tooltip": "An optional template to influence style of generation, created by the PixVerse Template node."
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         image: torch.Tensor,
         prompt: str,
         quality: str,
@@ -314,13 +312,11 @@ class PixverseImageToVideoNode(comfy_io.ComfyNode):
         seed,
         negative_prompt: str = None,
         pixverse_template: int = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: Optional[str] = None,
+        **kwargs,
+    ):
         validate_string(prompt, strip_whitespace=False)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        img_id = await upload_image_to_pixverse(image, auth_kwargs=auth)
+        img_id = await upload_image_to_pixverse(image, auth_kwargs=kwargs)
 
         # 1080p is limited to 5 seconds duration
         # only normal motion_mode supported for 1080p or for non-5 second duration
@@ -347,7 +343,7 @@ class PixverseImageToVideoNode(comfy_io.ComfyNode):
                 template_id=pixverse_template,
                 seed=seed,
             ),
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         response_api = await operation.execute()
 
@@ -368,8 +364,8 @@ class PixverseImageToVideoNode(comfy_io.ComfyNode):
                 PixverseStatus.deleted,
             ],
             status_extractor=lambda x: x.Resp.status,
-            auth_kwargs=auth,
-            node_id=cls.hidden.unique_id,
+            auth_kwargs=kwargs,
+            node_id=unique_id,
             result_url_extractor=get_video_url_from_response,
             estimated_duration=AVERAGE_DURATION_I2V,
         )
@@ -377,71 +373,72 @@ class PixverseImageToVideoNode(comfy_io.ComfyNode):
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.Resp.url) as vid_response:
-                return comfy_io.NodeOutput(VideoFromFile(BytesIO(await vid_response.content.read())))
+                return (VideoFromFile(BytesIO(await vid_response.content.read())),)
 
 
-class PixverseTransitionVideoNode(comfy_io.ComfyNode):
+class PixverseTransitionVideoNode(ComfyNodeABC):
     """
     Generates videos based on prompt and output_size.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="PixverseTransitionVideoNode",
-            display_name="PixVerse Transition Video",
-            category="api node/video/PixVerse",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("first_frame"),
-                comfy_io.Image.Input("last_frame"),
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt for the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "quality",
-                    options=PixverseQuality,
-                    default=PixverseQuality.res_540p,
-                ),
-                comfy_io.Combo.Input(
-                    "duration_seconds",
-                    options=PixverseDuration,
-                ),
-                comfy_io.Combo.Input(
-                    "motion_mode",
-                    options=PixverseMotionMode,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=2147483647,
-                    control_after_generate=True,
-                    tooltip="Seed for video generation.",
-                ),
-                comfy_io.String.Input(
-                    "negative_prompt",
-                    default="",
-                    multiline=True,
-                    tooltip="An optional text description of undesired elements on an image.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    RETURN_TYPES = (IO.VIDEO,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/video/PixVerse"
 
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "first_frame": (IO.IMAGE,),
+                "last_frame": (IO.IMAGE,),
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the video generation",
+                    },
+                ),
+                "quality": (
+                    [resolution.value for resolution in PixverseQuality],
+                    {
+                        "default": PixverseQuality.res_540p,
+                    },
+                ),
+                "duration_seconds": ([dur.value for dur in PixverseDuration],),
+                "motion_mode": ([mode.value for mode in PixverseMotionMode],),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2147483647,
+                        "control_after_generate": True,
+                        "tooltip": "Seed for video generation.",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    IO.STRING,
+                    {
+                        "default": "",
+                        "forceInput": True,
+                        "tooltip": "An optional text description of undesired elements on an image.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    async def api_call(
+        self,
         first_frame: torch.Tensor,
         last_frame: torch.Tensor,
         prompt: str,
@@ -450,14 +447,12 @@ class PixverseTransitionVideoNode(comfy_io.ComfyNode):
         motion_mode: str,
         seed,
         negative_prompt: str = None,
-    ) -> comfy_io.NodeOutput:
+        unique_id: Optional[str] = None,
+        **kwargs,
+    ):
         validate_string(prompt, strip_whitespace=False)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        first_frame_id = await upload_image_to_pixverse(first_frame, auth_kwargs=auth)
-        last_frame_id = await upload_image_to_pixverse(last_frame, auth_kwargs=auth)
+        first_frame_id = await upload_image_to_pixverse(first_frame, auth_kwargs=kwargs)
+        last_frame_id = await upload_image_to_pixverse(last_frame, auth_kwargs=kwargs)
 
         # 1080p is limited to 5 seconds duration
         # only normal motion_mode supported for 1080p or for non-5 second duration
@@ -484,7 +479,7 @@ class PixverseTransitionVideoNode(comfy_io.ComfyNode):
                 negative_prompt=negative_prompt if negative_prompt else None,
                 seed=seed,
             ),
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         response_api = await operation.execute()
 
@@ -505,8 +500,8 @@ class PixverseTransitionVideoNode(comfy_io.ComfyNode):
                 PixverseStatus.deleted,
             ],
             status_extractor=lambda x: x.Resp.status,
-            auth_kwargs=auth,
-            node_id=cls.hidden.unique_id,
+            auth_kwargs=kwargs,
+            node_id=unique_id,
             result_url_extractor=get_video_url_from_response,
             estimated_duration=AVERAGE_DURATION_T2V,
         )
@@ -514,19 +509,19 @@ class PixverseTransitionVideoNode(comfy_io.ComfyNode):
 
         async with aiohttp.ClientSession() as session:
             async with session.get(response_poll.Resp.url) as vid_response:
-                return comfy_io.NodeOutput(VideoFromFile(BytesIO(await vid_response.content.read())))
+                return (VideoFromFile(BytesIO(await vid_response.content.read())),)
 
 
-class PixVerseExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            PixverseTextToVideoNode,
-            PixverseImageToVideoNode,
-            PixverseTransitionVideoNode,
-            PixverseTemplateNode,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "PixverseTextToVideoNode": PixverseTextToVideoNode,
+    "PixverseImageToVideoNode": PixverseImageToVideoNode,
+    "PixverseTransitionVideoNode": PixverseTransitionVideoNode,
+    "PixverseTemplateNode": PixverseTemplateNode,
+}
 
-
-async def comfy_entrypoint() -> PixVerseExtension:
-    return PixVerseExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PixverseTextToVideoNode": "PixVerse Text to Video",
+    "PixverseImageToVideoNode": "PixVerse Image to Video",
+    "PixverseTransitionVideoNode": "PixVerse Transition Video",
+    "PixverseTemplateNode": "PixVerse Template",
+}

comfy_api_nodes/nodes_recraft.py

@@ -35,64 +35,57 @@ from server import PromptServer
 import torch
 from io import BytesIO
 from PIL import UnidentifiedImageError
-import aiohttp
 
 
 async def handle_recraft_file_request(
-    image: torch.Tensor,
-    path: str,
-    mask: torch.Tensor=None,
-    total_pixels=4096*4096,
-    timeout=1024,
-    request=None,
-    auth_kwargs: dict[str,str] = None,
-) -> list[BytesIO]:
+        image: torch.Tensor,
+        path: str,
+        mask: torch.Tensor=None,
+        total_pixels=4096*4096,
+        timeout=1024,
+        request=None,
+        auth_kwargs: dict[str,str] = None,
+    ) -> list[BytesIO]:
+        """
+        Handle sending common Recraft file-only request to get back file bytes.
+        """
+        if request is None:
+            request = EmptyRequest()
+
+        files = {
+            'image': tensor_to_bytesio(image, total_pixels=total_pixels).read()
+        }
+        if mask is not None:
+            files['mask'] = tensor_to_bytesio(mask, total_pixels=total_pixels).read()
+
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=path,
+                method=HttpMethod.POST,
+                request_model=type(request),
+                response_model=RecraftImageGenerationResponse,
+            ),
+            request=request,
+            files=files,
+            content_type="multipart/form-data",
+            auth_kwargs=auth_kwargs,
+            multipart_parser=recraft_multipart_parser,
+        )
+        response: RecraftImageGenerationResponse = await operation.execute()
+        all_bytesio = []
+        if response.image is not None:
+            all_bytesio.append(await download_url_to_bytesio(response.image.url, timeout=timeout))
+        else:
+            for data in response.data:
+                all_bytesio.append(await download_url_to_bytesio(data.url, timeout=timeout))
+
+        return all_bytesio
+
+
+def recraft_multipart_parser(data, parent_key=None, formatter: callable=None, converted_to_check: list[list]=None, is_list=False) -> dict:
     """
-    Handle sending common Recraft file-only request to get back file bytes.
-    """
-    if request is None:
-        request = EmptyRequest()
-
-    files = {
-        'image': tensor_to_bytesio(image, total_pixels=total_pixels).read()
-    }
-    if mask is not None:
-        files['mask'] = tensor_to_bytesio(mask, total_pixels=total_pixels).read()
-
-    operation = SynchronousOperation(
-        endpoint=ApiEndpoint(
-            path=path,
-            method=HttpMethod.POST,
-            request_model=type(request),
-            response_model=RecraftImageGenerationResponse,
-        ),
-        request=request,
-        files=files,
-        content_type="multipart/form-data",
-        auth_kwargs=auth_kwargs,
-        multipart_parser=recraft_multipart_parser,
-    )
-    response: RecraftImageGenerationResponse = await operation.execute()
-    all_bytesio = []
-    if response.image is not None:
-        all_bytesio.append(await download_url_to_bytesio(response.image.url, timeout=timeout))
-    else:
-        for data in response.data:
-            all_bytesio.append(await download_url_to_bytesio(data.url, timeout=timeout))
-
-    return all_bytesio
-
-
-def recraft_multipart_parser(
-    data,
-    parent_key=None,
-    formatter: callable = None,
-    converted_to_check: list[list] = None,
-    is_list: bool = False,
-    return_mode: str = "formdata"  # "dict" | "formdata"
-) -> dict | aiohttp.FormData:
-    """
-    Formats data such that multipart/form-data will work with aiohttp library when both files and data are present.
+    Formats data such that multipart/form-data will work with requests library
+    when both files and data are present.
 
     The OpenAI client that Recraft uses has a bizarre way of serializing lists:
 
@@ -110,23 +103,23 @@ def recraft_multipart_parser(
     # Modification of a function that handled a different type of multipart parsing, big ups:
     # https://gist.github.com/kazqvaizer/4cebebe5db654a414132809f9f88067b
 
-    def handle_converted_lists(item, parent_key, lists_to_check=tuple[list]):
+    def handle_converted_lists(data, parent_key, lists_to_check=tuple[list]):
         # if list already exists exists, just extend list with data
         for check_list in lists_to_check:
             for conv_tuple in check_list:
-                if conv_tuple[0] == parent_key and isinstance(conv_tuple[1], list):
-                    conv_tuple[1].append(formatter(item))
+                if conv_tuple[0] == parent_key and type(conv_tuple[1]) is list:
+                    conv_tuple[1].append(formatter(data))
                     return True
         return False
 
     if converted_to_check is None:
         converted_to_check = []
 
-    effective_mode = return_mode if parent_key is None else "dict"
+
     if formatter is None:
         formatter = lambda v: v  # Multipart representation of value
 
-    if not isinstance(data, dict):
+    if type(data) is not dict:
         # if list already exists exists, just extend list with data
         added = handle_converted_lists(data, parent_key, converted_to_check)
         if added:
@@ -143,24 +136,15 @@ def recraft_multipart_parser(
 
     for key, value in data.items():
         current_key = key if parent_key is None else f"{parent_key}[{key}]"
-        if isinstance(value, dict):
+        if type(value) is dict:
             converted.extend(recraft_multipart_parser(value, current_key, formatter, next_check).items())
-        elif isinstance(value, list):
+        elif type(value) is list:
             for ind, list_value in enumerate(value):
                 iter_key = f"{current_key}[]"
                 converted.extend(recraft_multipart_parser(list_value, iter_key, formatter, next_check, is_list=True).items())
         else:
             converted.append((current_key, formatter(value)))
 
-    if effective_mode == "formdata":
-        fd = aiohttp.FormData()
-        for k, v in dict(converted).items():
-            if isinstance(v, list):
-                for item in v:
-                    fd.add_field(k, str(item))
-            else:
-                fd.add_field(k, str(v))
-        return fd
     return dict(converted)
 
 

comfy_api_nodes/nodes_rodin.py

@@ -7,15 +7,15 @@ Rodin API docs: https://developer.hyper3d.ai/
 
 from __future__ import annotations
 from inspect import cleandoc
+from comfy.comfy_types.node_typing import IO
 import folder_paths as comfy_paths
 import aiohttp
 import os
+import datetime
 import asyncio
+import io
 import logging
 import math
-from typing import Optional
-from io import BytesIO
-from typing_extensions import override
 from PIL import Image
 from comfy_api_nodes.apis.rodin_api import (
     Rodin3DGenerateRequest,
@@ -32,548 +32,444 @@ from comfy_api_nodes.apis.client import (
     SynchronousOperation,
     PollingOperation,
 )
-from comfy_api.latest import ComfyExtension, io as comfy_io
 
 
-COMMON_PARAMETERS = [
-    comfy_io.Int.Input(
-        "Seed",
-        default=0,
-        min=0,
-        max=65535,
-        display_mode=comfy_io.NumberDisplay.number,
-        optional=True,
+COMMON_PARAMETERS = {
+    "Seed": (
+        IO.INT,
+        {
+            "default":0,
+            "min":0,
+            "max":65535,
+            "display":"number"
+        }
     ),
-    comfy_io.Combo.Input("Material_Type", options=["PBR", "Shaded"], default="PBR", optional=True),
-    comfy_io.Combo.Input(
-        "Polygon_count",
-        options=["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "200K-Triangle"],
-        default="18K-Quad",
-        optional=True,
+    "Material_Type": (
+        IO.COMBO,
+        {
+            "options": ["PBR", "Shaded"],
+            "default": "PBR"
+        }
     ),
-]
+    "Polygon_count": (
+        IO.COMBO,
+        {
+            "options": ["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "200K-Triangle"],
+            "default": "18K-Quad"
+        }
+    )
+}
+
+def create_task_error(response: Rodin3DGenerateResponse):
+    """Check if the response has error"""
+    return hasattr(response, "error")
 
 
-def get_quality_mode(poly_count):
-    polycount = poly_count.split("-")
-    poly = polycount[1]
-    count = polycount[0]
-    if poly == "Triangle":
-        mesh_mode = "Raw"
-    elif poly == "Quad":
-        mesh_mode = "Quad"
-    else:
-        mesh_mode = "Quad"
-
-    if count == "4K":
-        quality_override = 4000
-    elif count == "8K":
-        quality_override = 8000
-    elif count == "18K":
-        quality_override = 18000
-    elif count == "50K":
-        quality_override = 50000
-    elif count == "2K":
-        quality_override = 2000
-    elif count == "20K":
-        quality_override = 20000
-    elif count == "150K":
-        quality_override = 150000
-    elif count == "500K":
-        quality_override = 500000
-    else:
-        quality_override = 18000
-
-    return mesh_mode, quality_override
-
-
-def tensor_to_filelike(tensor, max_pixels: int = 2048*2048):
+class Rodin3DAPI:
     """
-    Converts a PyTorch tensor to a file-like object.
-
-    Args:
-    - tensor (torch.Tensor): A tensor representing an image of shape (H, W, C)
-      where C is the number of channels (3 for RGB), H is height, and W is width.
-
-    Returns:
-    - io.BytesIO: A file-like object containing the image data.
+    Generate 3D Assets using Rodin API
     """
-    array = tensor.cpu().numpy()
-    array = (array * 255).astype('uint8')
-    image = Image.fromarray(array, 'RGB')
+    RETURN_TYPES = (IO.STRING,)
+    RETURN_NAMES = ("3D Model Path",)
+    CATEGORY = "api node/3d/Rodin"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "api_call"
+    API_NODE = True
 
-    original_width, original_height = image.size
-    original_pixels = original_width * original_height
-    if original_pixels > max_pixels:
-        scale = math.sqrt(max_pixels / original_pixels)
-        new_width = int(original_width * scale)
-        new_height = int(original_height * scale)
-    else:
-        new_width, new_height = original_width, original_height
+    def tensor_to_filelike(self, tensor, max_pixels: int = 2048*2048):
+        """
+        Converts a PyTorch tensor to a file-like object.
 
-    if new_width != original_width or new_height != original_height:
-        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        Args:
+        - tensor (torch.Tensor): A tensor representing an image of shape (H, W, C)
+          where C is the number of channels (3 for RGB), H is height, and W is width.
 
-    img_byte_arr = BytesIO()
-    image.save(img_byte_arr, format='PNG')  # PNG is used for lossless compression
-    img_byte_arr.seek(0)
-    return img_byte_arr
+        Returns:
+        - io.BytesIO: A file-like object containing the image data.
+        """
+        array = tensor.cpu().numpy()
+        array = (array * 255).astype('uint8')
+        image = Image.fromarray(array, 'RGB')
 
+        original_width, original_height = image.size
+        original_pixels = original_width * original_height
+        if original_pixels > max_pixels:
+            scale = math.sqrt(max_pixels / original_pixels)
+            new_width = int(original_width * scale)
+            new_height = int(original_height * scale)
+        else:
+            new_width, new_height = original_width, original_height
 
-async def create_generate_task(
-    images=None,
-    seed=1,
-    material="PBR",
-    quality_override=18000,
-    tier="Regular",
-    mesh_mode="Quad",
-    TAPose = False,
-    auth_kwargs: Optional[dict[str, str]] = None,
-):
-    if images is None:
-        raise Exception("Rodin 3D generate requires at least 1 image.")
-    if len(images) > 5:
-        raise Exception("Rodin 3D generate requires up to 5 image.")
+        if new_width != original_width or new_height != original_height:
+            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
 
-    path = "/proxy/rodin/api/v2/rodin"
-    operation = SynchronousOperation(
-        endpoint=ApiEndpoint(
-            path=path,
-            method=HttpMethod.POST,
-            request_model=Rodin3DGenerateRequest,
-            response_model=Rodin3DGenerateResponse,
-        ),
-        request=Rodin3DGenerateRequest(
-            seed=seed,
-            tier=tier,
-            material=material,
-            quality_override=quality_override,
-            mesh_mode=mesh_mode,
-            TAPose=TAPose,
-        ),
-        files=[
-            (
-                "images",
-                open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image)
-            )
-            for image in images if image is not None
-        ],
-        content_type="multipart/form-data",
-        auth_kwargs=auth_kwargs,
-    )
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')  # PNG is used for lossless compression
+        img_byte_arr.seek(0)
+        return img_byte_arr
 
-    response = await operation.execute()
+    def check_rodin_status(self, response: Rodin3DCheckStatusResponse) -> str:
+        has_failed = any(job.status == JobStatus.Failed for job in response.jobs)
+        all_done = all(job.status == JobStatus.Done for job in response.jobs)
+        status_list = [str(job.status) for job in response.jobs]
+        logging.info(f"[ Rodin3D API - CheckStatus ] Generate Status: {status_list}")
+        if has_failed:
+            logging.error(f"[ Rodin3D API - CheckStatus ] Generate Failed: {status_list}, Please try again.")
+            raise Exception("[ Rodin3D API ] Generate Failed, Please Try again.")
+        elif all_done:
+            return "DONE"
+        else:
+            return "Generating"
 
-    if hasattr(response, "error"):
-        error_message = f"Rodin3D Create 3D generate Task Failed. Message: {response.message}, error: {response.error}"
-        logging.error(error_message)
-        raise Exception(error_message)
+    async def create_generate_task(self, images=None, seed=1, material="PBR", quality="medium", tier="Regular", mesh_mode="Quad", **kwargs):
+        if images is None:
+            raise Exception("Rodin 3D generate requires at least 1 image.")
+        if len(images) >= 5:
+            raise Exception("Rodin 3D generate requires up to 5 image.")
 
-    logging.info("[ Rodin3D API - Submit Jobs ] Submit Generate Task Success!")
-    subscription_key = response.jobs.subscription_key
-    task_uuid = response.uuid
-    logging.info("[ Rodin3D API - Submit Jobs ] UUID: %s", task_uuid)
-    return task_uuid, subscription_key
-
-
-def check_rodin_status(response: Rodin3DCheckStatusResponse) -> str:
-    all_done = all(job.status == JobStatus.Done for job in response.jobs)
-    status_list = [str(job.status) for job in response.jobs]
-    logging.info("[ Rodin3D API - CheckStatus ] Generate Status: %s", status_list)
-    if any(job.status == JobStatus.Failed for job in response.jobs):
-        logging.error("[ Rodin3D API - CheckStatus ] Generate Failed: %s, Please try again.", status_list)
-        raise Exception("[ Rodin3D API ] Generate Failed, Please Try again.")
-    if all_done:
-        return "DONE"
-    return "Generating"
-
-
-async def poll_for_task_status(
-    subscription_key, auth_kwargs: Optional[dict[str, str]] = None,
-) -> Rodin3DCheckStatusResponse:
-    poll_operation = PollingOperation(
-        poll_endpoint=ApiEndpoint(
-            path="/proxy/rodin/api/v2/status",
-            method=HttpMethod.POST,
-            request_model=Rodin3DCheckStatusRequest,
-            response_model=Rodin3DCheckStatusResponse,
-        ),
-        request=Rodin3DCheckStatusRequest(subscription_key=subscription_key),
-        completed_statuses=["DONE"],
-        failed_statuses=["FAILED"],
-        status_extractor=check_rodin_status,
-        poll_interval=3.0,
-        auth_kwargs=auth_kwargs,
-    )
-    logging.info("[ Rodin3D API - CheckStatus ] Generate Start!")
-    return await poll_operation.execute()
-
-
-async def get_rodin_download_list(uuid, auth_kwargs: Optional[dict[str, str]] = None) -> Rodin3DDownloadResponse:
-    logging.info("[ Rodin3D API - Downloading ] Generate Successfully!")
-    operation = SynchronousOperation(
-        endpoint=ApiEndpoint(
-            path="/proxy/rodin/api/v2/download",
-            method=HttpMethod.POST,
-            request_model=Rodin3DDownloadRequest,
-            response_model=Rodin3DDownloadResponse,
-        ),
-        request=Rodin3DDownloadRequest(task_uuid=uuid),
-        auth_kwargs=auth_kwargs,
-    )
-    return await operation.execute()
-
-
-async def download_files(url_list, task_uuid):
-    save_path = os.path.join(comfy_paths.get_output_directory(), f"Rodin3D_{task_uuid}")
-    os.makedirs(save_path, exist_ok=True)
-    model_file_path = None
-    async with aiohttp.ClientSession() as session:
-        for i in url_list.list:
-            url = i.url
-            file_name = i.name
-            file_path = os.path.join(save_path, file_name)
-            if file_path.endswith(".glb"):
-                model_file_path = file_path
-            logging.info("[ Rodin3D API - download_files ] Downloading file: %s", file_path)
-            max_retries = 5
-            for attempt in range(max_retries):
-                try:
-                    async with session.get(url) as resp:
-                        resp.raise_for_status()
-                        with open(file_path, "wb") as f:
-                            async for chunk in resp.content.iter_chunked(32 * 1024):
-                                f.write(chunk)
-                    break
-                except Exception as e:
-                    logging.info("[ Rodin3D API - download_files ] Error downloading %s:%s", file_path, str(e))
-                    if attempt < max_retries - 1:
-                        logging.info("Retrying...")
-                        await asyncio.sleep(2)
-                    else:
-                        logging.info(
-                            "[ Rodin3D API - download_files ] Failed to download %s after %s attempts.",
-                            file_path,
-                            max_retries,
-                        )
-    return model_file_path
-
-
-class Rodin3D_Regular(comfy_io.ComfyNode):
-    """Generate 3D Assets using Rodin API"""
-
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="Rodin3D_Regular",
-            display_name="Rodin 3D Generate - Regular Generate",
-            category="api node/3d/Rodin",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("Images"),
-                *COMMON_PARAMETERS,
+        path = "/proxy/rodin/api/v2/rodin"
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=path,
+                method=HttpMethod.POST,
+                request_model=Rodin3DGenerateRequest,
+                response_model=Rodin3DGenerateResponse,
+            ),
+            request=Rodin3DGenerateRequest(
+                seed=seed,
+                tier=tier,
+                material=material,
+                quality=quality,
+                mesh_mode=mesh_mode
+            ),
+            files=[
+                (
+                    "images",
+                    open(image, "rb") if isinstance(image, str) else self.tensor_to_filelike(image)
+                )
+                for image in images if image is not None
             ],
-            outputs=[comfy_io.String.Output(display_name="3D Model Path")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-            ],
-            is_api_node=True,
+            content_type = "multipart/form-data",
+            auth_kwargs=kwargs,
         )
 
+        response = await operation.execute()
+
+        if create_task_error(response):
+            error_message = f"Rodin3D Create 3D generate Task Failed. Message: {response.message}, error: {response.error}"
+            logging.error(error_message)
+            raise Exception(error_message)
+
+        logging.info("[ Rodin3D API - Submit Jobs ] Submit Generate Task Success!")
+        subscription_key = response.jobs.subscription_key
+        task_uuid = response.uuid
+        logging.info(f"[ Rodin3D API - Submit Jobs ] UUID: {task_uuid}")
+        return task_uuid, subscription_key
+
+    async def poll_for_task_status(self, subscription_key, **kwargs) -> Rodin3DCheckStatusResponse:
+
+        path = "/proxy/rodin/api/v2/status"
+
+        poll_operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path = path,
+                method=HttpMethod.POST,
+                request_model=Rodin3DCheckStatusRequest,
+                response_model=Rodin3DCheckStatusResponse,
+            ),
+            request=Rodin3DCheckStatusRequest(
+                subscription_key = subscription_key
+            ),
+            completed_statuses=["DONE"],
+            failed_statuses=["FAILED"],
+            status_extractor=self.check_rodin_status,
+            poll_interval=3.0,
+            auth_kwargs=kwargs,
+        )
+
+        logging.info("[ Rodin3D API - CheckStatus ] Generate Start!")
+
+        return await poll_operation.execute()
+
+    async def get_rodin_download_list(self, uuid, **kwargs) -> Rodin3DDownloadResponse:
+        logging.info("[ Rodin3D API - Downloading ] Generate Successfully!")
+
+        path = "/proxy/rodin/api/v2/download"
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=path,
+                method=HttpMethod.POST,
+                request_model=Rodin3DDownloadRequest,
+                response_model=Rodin3DDownloadResponse,
+            ),
+            request=Rodin3DDownloadRequest(
+                task_uuid=uuid
+            ),
+            auth_kwargs=kwargs
+        )
+
+        return await operation.execute()
+
+    def get_quality_mode(self, poly_count):
+        if poly_count == "200K-Triangle":
+            mesh_mode = "Raw"
+            quality = "medium"
+        else:
+            mesh_mode = "Quad"
+            if poly_count == "4K-Quad":
+                quality = "extra-low"
+            elif poly_count == "8K-Quad":
+                quality = "low"
+            elif poly_count == "18K-Quad":
+                quality = "medium"
+            elif poly_count == "50K-Quad":
+                quality = "high"
+            else:
+                quality = "medium"
+
+        return mesh_mode, quality
+
+    async def download_files(self, url_list):
+        save_path = os.path.join(comfy_paths.get_output_directory(), "Rodin3D", datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
+        os.makedirs(save_path, exist_ok=True)
+        model_file_path = None
+        async with aiohttp.ClientSession() as session:
+            for i in url_list.list:
+                url = i.url
+                file_name = i.name
+                file_path = os.path.join(save_path, file_name)
+                if file_path.endswith(".glb"):
+                    model_file_path = file_path
+                logging.info(f"[ Rodin3D API - download_files ] Downloading file: {file_path}")
+                max_retries = 5
+                for attempt in range(max_retries):
+                    try:
+                        async with session.get(url) as resp:
+                            resp.raise_for_status()
+                            with open(file_path, "wb") as f:
+                                async for chunk in resp.content.iter_chunked(32 * 1024):
+                                    f.write(chunk)
+                        break
+                    except Exception as e:
+                        logging.info(f"[ Rodin3D API - download_files ] Error downloading {file_path}:{e}")
+                        if attempt < max_retries - 1:
+                            logging.info("Retrying...")
+                            await asyncio.sleep(2)
+                        else:
+                            logging.info(
+                                "[ Rodin3D API - download_files ] Failed to download %s after %s attempts.",
+                                file_path,
+                                max_retries,
+                            )
+
+        return model_file_path
+
+
+class Rodin3D_Regular(Rodin3DAPI):
     @classmethod
-    async def execute(
-        cls,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Images":
+                (
+                    IO.IMAGE,
+                    {
+                        "forceInput":True,
+                    }
+                )
+            },
+            "optional": {
+                **COMMON_PARAMETERS
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+            },
+        }
+
+    async def api_call(
+        self,
         Images,
         Seed,
         Material_Type,
         Polygon_count,
-    ) -> comfy_io.NodeOutput:
+        **kwargs
+    ):
         tier = "Regular"
         num_images = Images.shape[0]
         m_images = []
         for i in range(num_images):
             m_images.append(Images[i])
-        mesh_mode, quality_override = get_quality_mode(Polygon_count)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        mesh_mode, quality = self.get_quality_mode(Polygon_count)
+        task_uuid, subscription_key = await self.create_generate_task(images=m_images, seed=Seed, material=Material_Type,
+                                                                quality=quality, tier=tier, mesh_mode=mesh_mode,
+                                                                **kwargs)
+        await self.poll_for_task_status(subscription_key, **kwargs)
+        download_list = await self.get_rodin_download_list(task_uuid, **kwargs)
+        model = await self.download_files(download_list)
+
+        return (model,)
+
+
+class Rodin3D_Detail(Rodin3DAPI):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Images":
+                (
+                    IO.IMAGE,
+                    {
+                        "forceInput":True,
+                    }
+                )
+            },
+            "optional": {
+                **COMMON_PARAMETERS
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+            },
         }
-        task_uuid, subscription_key = await create_generate_task(
-            images=m_images,
-            seed=Seed,
-            material=Material_Type,
-            quality_override=quality_override,
-            tier=tier,
-            mesh_mode=mesh_mode,
-            auth_kwargs=auth,
-        )
-        await poll_for_task_status(subscription_key, auth_kwargs=auth)
-        download_list = await get_rodin_download_list(task_uuid, auth_kwargs=auth)
-        model = await download_files(download_list, task_uuid)
 
-        return comfy_io.NodeOutput(model)
-
-
-class Rodin3D_Detail(comfy_io.ComfyNode):
-    """Generate 3D Assets using Rodin API"""
-
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="Rodin3D_Detail",
-            display_name="Rodin 3D Generate - Detail Generate",
-            category="api node/3d/Rodin",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("Images"),
-                *COMMON_PARAMETERS,
-            ],
-            outputs=[comfy_io.String.Output(display_name="3D Model Path")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
+    async def api_call(
+        self,
         Images,
         Seed,
         Material_Type,
         Polygon_count,
-    ) -> comfy_io.NodeOutput:
+        **kwargs
+    ):
         tier = "Detail"
         num_images = Images.shape[0]
         m_images = []
         for i in range(num_images):
             m_images.append(Images[i])
-        mesh_mode, quality_override = get_quality_mode(Polygon_count)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        mesh_mode, quality = self.get_quality_mode(Polygon_count)
+        task_uuid, subscription_key = await self.create_generate_task(images=m_images, seed=Seed, material=Material_Type,
+                                                                quality=quality, tier=tier, mesh_mode=mesh_mode,
+                                                                **kwargs)
+        await self.poll_for_task_status(subscription_key, **kwargs)
+        download_list = await self.get_rodin_download_list(task_uuid, **kwargs)
+        model = await self.download_files(download_list)
+
+        return (model,)
+
+
+class Rodin3D_Smooth(Rodin3DAPI):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Images":
+                (
+                    IO.IMAGE,
+                    {
+                        "forceInput":True,
+                    }
+                )
+            },
+            "optional": {
+                **COMMON_PARAMETERS
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+            },
         }
-        task_uuid, subscription_key = await create_generate_task(
-            images=m_images,
-            seed=Seed,
-            material=Material_Type,
-            quality_override=quality_override,
-            tier=tier,
-            mesh_mode=mesh_mode,
-            auth_kwargs=auth,
-        )
-        await poll_for_task_status(subscription_key, auth_kwargs=auth)
-        download_list = await get_rodin_download_list(task_uuid, auth_kwargs=auth)
-        model = await download_files(download_list, task_uuid)
 
-        return comfy_io.NodeOutput(model)
-
-
-class Rodin3D_Smooth(comfy_io.ComfyNode):
-    """Generate 3D Assets using Rodin API"""
-
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="Rodin3D_Smooth",
-            display_name="Rodin 3D Generate - Smooth Generate",
-            category="api node/3d/Rodin",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("Images"),
-                *COMMON_PARAMETERS,
-            ],
-            outputs=[comfy_io.String.Output(display_name="3D Model Path")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
+    async def api_call(
+        self,
         Images,
         Seed,
         Material_Type,
         Polygon_count,
-    ) -> comfy_io.NodeOutput:
+        **kwargs
+    ):
         tier = "Smooth"
         num_images = Images.shape[0]
         m_images = []
         for i in range(num_images):
             m_images.append(Images[i])
-        mesh_mode, quality_override = get_quality_mode(Polygon_count)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        mesh_mode, quality = self.get_quality_mode(Polygon_count)
+        task_uuid, subscription_key = await self.create_generate_task(images=m_images, seed=Seed, material=Material_Type,
+                                                                quality=quality, tier=tier, mesh_mode=mesh_mode,
+                                                                **kwargs)
+        await self.poll_for_task_status(subscription_key, **kwargs)
+        download_list = await self.get_rodin_download_list(task_uuid, **kwargs)
+        model = await self.download_files(download_list)
+
+        return (model,)
+
+
+class Rodin3D_Sketch(Rodin3DAPI):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Images":
+                (
+                    IO.IMAGE,
+                    {
+                        "forceInput":True,
+                    }
+                )
+            },
+            "optional": {
+                "Seed":
+                (
+                    IO.INT,
+                    {
+                        "default":0,
+                        "min":0,
+                        "max":65535,
+                        "display":"number"
+                    }
+                )
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+            },
         }
-        task_uuid, subscription_key = await create_generate_task(
-            images=m_images,
-            seed=Seed,
-            material=Material_Type,
-            quality_override=quality_override,
-            tier=tier,
-            mesh_mode=mesh_mode,
-            auth_kwargs=auth,
-        )
-        await poll_for_task_status(subscription_key, auth_kwargs=auth)
-        download_list = await get_rodin_download_list(task_uuid, auth_kwargs=auth)
-        model = await download_files(download_list, task_uuid)
 
-        return comfy_io.NodeOutput(model)
-
-
-class Rodin3D_Sketch(comfy_io.ComfyNode):
-    """Generate 3D Assets using Rodin API"""
-
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="Rodin3D_Sketch",
-            display_name="Rodin 3D Generate - Sketch Generate",
-            category="api node/3d/Rodin",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("Images"),
-                comfy_io.Int.Input(
-                    "Seed",
-                    default=0,
-                    min=0,
-                    max=65535,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.String.Output(display_name="3D Model Path")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
+    async def api_call(
+        self,
         Images,
         Seed,
-    ) -> comfy_io.NodeOutput:
+        **kwargs
+    ):
         tier = "Sketch"
         num_images = Images.shape[0]
         m_images = []
         for i in range(num_images):
             m_images.append(Images[i])
         material_type = "PBR"
-        quality_override = 18000
+        quality = "medium"
         mesh_mode = "Quad"
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        task_uuid, subscription_key = await create_generate_task(
-            images=m_images,
-            seed=Seed,
-            material=material_type,
-            quality_override=quality_override,
-            tier=tier,
-            mesh_mode=mesh_mode,
-            auth_kwargs=auth,
+        task_uuid, subscription_key = await self.create_generate_task(
+            images=m_images, seed=Seed, material=material_type, quality=quality, tier=tier, mesh_mode=mesh_mode, **kwargs
         )
-        await poll_for_task_status(subscription_key, auth_kwargs=auth)
-        download_list = await get_rodin_download_list(task_uuid, auth_kwargs=auth)
-        model = await download_files(download_list, task_uuid)
+        await self.poll_for_task_status(subscription_key, **kwargs)
+        download_list = await self.get_rodin_download_list(task_uuid, **kwargs)
+        model = await self.download_files(download_list)
 
-        return comfy_io.NodeOutput(model)
+        return (model,)
 
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "Rodin3D_Regular": Rodin3D_Regular,
+    "Rodin3D_Detail": Rodin3D_Detail,
+    "Rodin3D_Smooth": Rodin3D_Smooth,
+    "Rodin3D_Sketch": Rodin3D_Sketch,
+}
 
-class Rodin3D_Gen2(comfy_io.ComfyNode):
-    """Generate 3D Assets using Rodin API"""
-
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="Rodin3D_Gen2",
-            display_name="Rodin 3D Generate - Gen-2 Generate",
-            category="api node/3d/Rodin",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input("Images"),
-                comfy_io.Int.Input(
-                    "Seed",
-                    default=0,
-                    min=0,
-                    max=65535,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    optional=True,
-                ),
-                comfy_io.Combo.Input("Material_Type", options=["PBR", "Shaded"], default="PBR", optional=True),
-                comfy_io.Combo.Input(
-                    "Polygon_count",
-                    options=["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "2K-Triangle", "20K-Triangle", "150K-Triangle", "500K-Triangle"],
-                    default="500K-Triangle",
-                    optional=True,
-                ),
-                comfy_io.Boolean.Input("TAPose", default=False),
-            ],
-            outputs=[comfy_io.String.Output(display_name="3D Model Path")],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        Images,
-        Seed,
-        Material_Type,
-        Polygon_count,
-        TAPose,
-    ) -> comfy_io.NodeOutput:
-        tier = "Gen-2"
-        num_images = Images.shape[0]
-        m_images = []
-        for i in range(num_images):
-            m_images.append(Images[i])
-        mesh_mode, quality_override = get_quality_mode(Polygon_count)
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        task_uuid, subscription_key = await create_generate_task(
-            images=m_images,
-            seed=Seed,
-            material=Material_Type,
-            quality_override=quality_override,
-            tier=tier,
-            mesh_mode=mesh_mode,
-            TAPose=TAPose,
-            auth_kwargs=auth,
-        )
-        await poll_for_task_status(subscription_key, auth_kwargs=auth)
-        download_list = await get_rodin_download_list(task_uuid, auth_kwargs=auth)
-        model = await download_files(download_list, task_uuid)
-
-        return comfy_io.NodeOutput(model)
-
-
-class Rodin3DExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            Rodin3D_Regular,
-            Rodin3D_Detail,
-            Rodin3D_Smooth,
-            Rodin3D_Sketch,
-            Rodin3D_Gen2,
-        ]
-
-
-async def comfy_entrypoint() -> Rodin3DExtension:
-    return Rodin3DExtension()
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Rodin3D_Regular": "Rodin 3D Generate - Regular Generate",
+    "Rodin3D_Detail": "Rodin 3D Generate - Detail Generate",
+    "Rodin3D_Smooth": "Rodin 3D Generate - Smooth Generate",
+    "Rodin3D_Sketch": "Rodin 3D Generate - Sketch Generate",
+}

comfy_api_nodes/nodes_runway.py

@@ -200,11 +200,11 @@ class RunwayImageToVideoNodeGen3a(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "duration",
-                    options=Duration,
+                    options=[model.value for model in Duration],
                 ),
                 comfy_io.Combo.Input(
                     "ratio",
-                    options=RunwayGen3aAspectRatio,
+                    options=[model.value for model in RunwayGen3aAspectRatio],
                 ),
                 comfy_io.Int.Input(
                     "seed",
@@ -300,11 +300,11 @@ class RunwayImageToVideoNodeGen4(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "duration",
-                    options=Duration,
+                    options=[model.value for model in Duration],
                 ),
                 comfy_io.Combo.Input(
                     "ratio",
-                    options=RunwayGen4TurboAspectRatio,
+                    options=[model.value for model in RunwayGen4TurboAspectRatio],
                 ),
                 comfy_io.Int.Input(
                     "seed",
@@ -408,11 +408,11 @@ class RunwayFirstLastFrameNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "duration",
-                    options=Duration,
+                    options=[model.value for model in Duration],
                 ),
                 comfy_io.Combo.Input(
                     "ratio",
-                    options=RunwayGen3aAspectRatio,
+                    options=[model.value for model in RunwayGen3aAspectRatio],
                 ),
                 comfy_io.Int.Input(
                     "seed",

comfy_api_nodes/nodes_sora.py

@@ -1,175 +0,0 @@
-from typing import Optional
-from typing_extensions import override
-
-import torch
-from pydantic import BaseModel, Field
-from comfy_api.latest import ComfyExtension, io as comfy_io
-from comfy_api_nodes.apis.client import (
-    ApiEndpoint,
-    HttpMethod,
-    SynchronousOperation,
-    PollingOperation,
-    EmptyRequest,
-)
-from comfy_api_nodes.util.validation_utils import get_number_of_images
-
-from comfy_api_nodes.apinode_utils import (
-    download_url_to_video_output,
-    tensor_to_bytesio,
-)
-
-class Sora2GenerationRequest(BaseModel):
-    prompt: str = Field(...)
-    model: str = Field(...)
-    seconds: str = Field(...)
-    size: str = Field(...)
-
-
-class Sora2GenerationResponse(BaseModel):
-    id: str = Field(...)
-    error: Optional[dict] = Field(None)
-    status: Optional[str] = Field(None)
-
-
-class OpenAIVideoSora2(comfy_io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return comfy_io.Schema(
-            node_id="OpenAIVideoSora2",
-            display_name="OpenAI Sora - Video",
-            category="api node/video/Sora",
-            description="OpenAI video and audio generation.",
-            inputs=[
-                comfy_io.Combo.Input(
-                    "model",
-                    options=["sora-2", "sora-2-pro"],
-                    default="sora-2",
-                ),
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Guiding text; may be empty if an input image is present.",
-                ),
-                comfy_io.Combo.Input(
-                    "size",
-                    options=[
-                        "720x1280",
-                        "1280x720",
-                        "1024x1792",
-                        "1792x1024",
-                    ],
-                    default="1280x720",
-                ),
-                comfy_io.Combo.Input(
-                    "duration",
-                    options=[4, 8, 12],
-                    default=8,
-                ),
-                comfy_io.Image.Input(
-                    "image",
-                    optional=True,
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=2147483647,
-                    step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    control_after_generate=True,
-                    optional=True,
-                    tooltip="Seed to determine if node should re-run; "
-                            "actual results are nondeterministic regardless of seed.",
-                ),
-            ],
-            outputs=[
-                comfy_io.Video.Output(),
-            ],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        model: str,
-        prompt: str,
-        size: str = "1280x720",
-        duration: int = 8,
-        seed: int = 0,
-        image: Optional[torch.Tensor] = None,
-    ):
-        if model == "sora-2" and size not in ("720x1280", "1280x720"):
-            raise ValueError("Invalid size for sora-2 model, only 720x1280 and 1280x720 are supported.")
-        files_input = None
-        if image is not None:
-            if get_number_of_images(image) != 1:
-                raise ValueError("Currently only one input image is supported.")
-            files_input = {"input_reference": ("image.png", tensor_to_bytesio(image), "image/png")}
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
-        payload = Sora2GenerationRequest(
-            model=model,
-            prompt=prompt,
-            seconds=str(duration),
-            size=size,
-        )
-        initial_operation = SynchronousOperation(
-            endpoint=ApiEndpoint(
-                path="/proxy/openai/v1/videos",
-                method=HttpMethod.POST,
-                request_model=Sora2GenerationRequest,
-                response_model=Sora2GenerationResponse
-            ),
-            request=payload,
-            files=files_input,
-            auth_kwargs=auth,
-            content_type="multipart/form-data",
-        )
-        initial_response = await initial_operation.execute()
-        if initial_response.error:
-            raise Exception(initial_response.error.message)
-
-        model_time_multiplier = 1 if model == "sora-2" else 2
-        poll_operation = PollingOperation(
-            poll_endpoint=ApiEndpoint(
-                path=f"/proxy/openai/v1/videos/{initial_response.id}",
-                method=HttpMethod.GET,
-                request_model=EmptyRequest,
-                response_model=Sora2GenerationResponse
-            ),
-            completed_statuses=["completed"],
-            failed_statuses=["failed"],
-            status_extractor=lambda x: x.status,
-            auth_kwargs=auth,
-            poll_interval=8.0,
-            max_poll_attempts=160,
-            node_id=cls.hidden.unique_id,
-            estimated_duration=45 * (duration / 4) * model_time_multiplier,
-        )
-        await poll_operation.execute()
-        return comfy_io.NodeOutput(
-            await download_url_to_video_output(
-                f"/proxy/openai/v1/videos/{initial_response.id}/content",
-                auth_kwargs=auth,
-            )
-        )
-
-
-class OpenAISoraExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            OpenAIVideoSora2,
-        ]
-
-
-async def comfy_entrypoint() -> OpenAISoraExtension:
-    return OpenAISoraExtension()

comfy_api_nodes/nodes_stability.py

@@ -82,8 +82,8 @@ class StabilityStableImageUltraNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "aspect_ratio",
-                    options=StabilityAspectRatio,
-                    default=StabilityAspectRatio.ratio_1_1,
+                    options=[x.value for x in StabilityAspectRatio],
+                    default=StabilityAspectRatio.ratio_1_1.value,
                     tooltip="Aspect ratio of generated image.",
                 ),
                 comfy_io.Combo.Input(
@@ -217,12 +217,12 @@ class StabilityStableImageSD_3_5Node(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "model",
-                    options=Stability_SD3_5_Model,
+                    options=[x.value for x in Stability_SD3_5_Model],
                 ),
                 comfy_io.Combo.Input(
                     "aspect_ratio",
-                    options=StabilityAspectRatio,
-                    default=StabilityAspectRatio.ratio_1_1,
+                    options=[x.value for x in StabilityAspectRatio],
+                    default=StabilityAspectRatio.ratio_1_1.value,
                     tooltip="Aspect ratio of generated image.",
                 ),
                 comfy_io.Combo.Input(

comfy_api_nodes/nodes_veo2.py

@@ -215,7 +215,7 @@ class VeoVideoGenerationNode(comfy_io.ComfyNode):
         initial_response = await initial_operation.execute()
         operation_name = initial_response.name
 
-        logging.info("Veo generation started with operation name: %s", operation_name)
+        logging.info(f"Veo generation started with operation name: {operation_name}")
 
         # Define status extractor function
         def status_extractor(response):

comfy_api_nodes/nodes_vidu.py

@@ -173,8 +173,8 @@ class ViduTextToVideoNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=VideoModelName,
-                    default=VideoModelName.vidu_q1,
+                    options=[model.value for model in VideoModelName],
+                    default=VideoModelName.vidu_q1.value,
                     tooltip="Model name",
                 ),
                 comfy_io.String.Input(
@@ -205,22 +205,22 @@ class ViduTextToVideoNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "aspect_ratio",
-                    options=AspectRatio,
-                    default=AspectRatio.r_16_9,
+                    options=[model.value for model in AspectRatio],
+                    default=AspectRatio.r_16_9.value,
                     tooltip="The aspect ratio of the output video",
                     optional=True,
                 ),
                 comfy_io.Combo.Input(
                     "resolution",
-                    options=Resolution,
-                    default=Resolution.r_1080p,
+                    options=[model.value for model in Resolution],
+                    default=Resolution.r_1080p.value,
                     tooltip="Supported values may vary by model & duration",
                     optional=True,
                 ),
                 comfy_io.Combo.Input(
                     "movement_amplitude",
-                    options=MovementAmplitude,
-                    default=MovementAmplitude.auto,
+                    options=[model.value for model in MovementAmplitude],
+                    default=MovementAmplitude.auto.value,
                     tooltip="The movement amplitude of objects in the frame",
                     optional=True,
                 ),
@@ -278,8 +278,8 @@ class ViduImageToVideoNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=VideoModelName,
-                    default=VideoModelName.vidu_q1,
+                    options=[model.value for model in VideoModelName],
+                    default=VideoModelName.vidu_q1.value,
                     tooltip="Model name",
                 ),
                 comfy_io.Image.Input(
@@ -316,14 +316,14 @@ class ViduImageToVideoNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "resolution",
-                    options=Resolution,
-                    default=Resolution.r_1080p,
+                    options=[model.value for model in Resolution],
+                    default=Resolution.r_1080p.value,
                     tooltip="Supported values may vary by model & duration",
                     optional=True,
                 ),
                 comfy_io.Combo.Input(
                     "movement_amplitude",
-                    options=MovementAmplitude,
+                    options=[model.value for model in MovementAmplitude],
                     default=MovementAmplitude.auto.value,
                     tooltip="The movement amplitude of objects in the frame",
                     optional=True,
@@ -388,8 +388,8 @@ class ViduReferenceVideoNode(comfy_io.ComfyNode):
             inputs=[
                 comfy_io.Combo.Input(
                     "model",
-                    options=VideoModelName,
-                    default=VideoModelName.vidu_q1,
+                    options=[model.value for model in VideoModelName],
+                    default=VideoModelName.vidu_q1.value,
                     tooltip="Model name",
                 ),
                 comfy_io.Image.Input(
@@ -424,8 +424,8 @@ class ViduReferenceVideoNode(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "aspect_ratio",
-                    options=AspectRatio,
-                    default=AspectRatio.r_16_9,
+                    options=[model.value for model in AspectRatio],
+                    default=AspectRatio.r_16_9.value,
                     tooltip="The aspect ratio of the output video",
                     optional=True,
                 ),

comfy_api_nodes/nodes_wan.py

@@ -28,12 +28,6 @@ class Text2ImageInputField(BaseModel):
     negative_prompt: Optional[str] = Field(None)
 
 
-class Image2ImageInputField(BaseModel):
-    prompt: str = Field(...)
-    negative_prompt: Optional[str] = Field(None)
-    images: list[str] = Field(..., min_length=1, max_length=2)
-
-
 class Text2VideoInputField(BaseModel):
     prompt: str = Field(...)
     negative_prompt: Optional[str] = Field(None)
@@ -55,13 +49,6 @@ class Txt2ImageParametersField(BaseModel):
     watermark: bool = Field(True)
 
 
-class Image2ImageParametersField(BaseModel):
-    size: Optional[str] = Field(None)
-    n: int = Field(1, description="Number of images to generate.")  # we support only value=1
-    seed: int = Field(..., ge=0, le=2147483647)
-    watermark: bool = Field(True)
-
-
 class Text2VideoParametersField(BaseModel):
     size: str = Field(...)
     seed: int = Field(..., ge=0, le=2147483647)
@@ -86,12 +73,6 @@ class Text2ImageTaskCreationRequest(BaseModel):
     parameters: Txt2ImageParametersField = Field(...)
 
 
-class Image2ImageTaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    input: Image2ImageInputField = Field(...)
-    parameters: Image2ImageParametersField = Field(...)
-
-
 class Text2VideoTaskCreationRequest(BaseModel):
     model: str = Field(...)
     input: Text2VideoInputField = Field(...)
@@ -154,12 +135,7 @@ async def process_task(
     url: str,
     request_model: Type[T],
     response_model: Type[R],
-    payload: Union[
-        Text2ImageTaskCreationRequest,
-        Image2ImageTaskCreationRequest,
-        Text2VideoTaskCreationRequest,
-        Image2VideoTaskCreationRequest,
-    ],
+    payload: Union[Text2ImageTaskCreationRequest, Text2VideoTaskCreationRequest, Image2VideoTaskCreationRequest],
     node_id: str,
     estimated_duration: int,
     poll_interval: int,
@@ -312,128 +288,6 @@ class WanTextToImageApi(comfy_io.ComfyNode):
         return comfy_io.NodeOutput(await download_url_to_image_tensor(str(response.output.results[0].url)))
 
 
-class WanImageToImageApi(comfy_io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return comfy_io.Schema(
-            node_id="WanImageToImageApi",
-            display_name="Wan Image to Image",
-            category="api node/image/Wan",
-            description="Generates an image from one or two input images and a text prompt. "
-                        "The output image is currently fixed at 1.6 MP; its aspect ratio matches the input image(s).",
-            inputs=[
-                comfy_io.Combo.Input(
-                    "model",
-                    options=["wan2.5-i2i-preview"],
-                    default="wan2.5-i2i-preview",
-                    tooltip="Model to use.",
-                ),
-                comfy_io.Image.Input(
-                    "image",
-                    tooltip="Single-image editing or multi-image fusion, maximum 2 images.",
-                ),
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Prompt used to describe the elements and visual features, supports English/Chinese.",
-                ),
-                comfy_io.String.Input(
-                    "negative_prompt",
-                    multiline=True,
-                    default="",
-                    tooltip="Negative text prompt to guide what to avoid.",
-                    optional=True,
-                ),
-                # redo this later as an optional combo of recommended resolutions
-                # comfy_io.Int.Input(
-                #     "width",
-                #     default=1280,
-                #     min=384,
-                #     max=1440,
-                #     step=16,
-                #     optional=True,
-                # ),
-                # comfy_io.Int.Input(
-                #     "height",
-                #     default=1280,
-                #     min=384,
-                #     max=1440,
-                #     step=16,
-                #     optional=True,
-                # ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=2147483647,
-                    step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    control_after_generate=True,
-                    tooltip="Seed to use for generation.",
-                    optional=True,
-                ),
-                comfy_io.Boolean.Input(
-                    "watermark",
-                    default=True,
-                    tooltip="Whether to add an \"AI generated\" watermark to the result.",
-                    optional=True,
-                ),
-            ],
-            outputs=[
-                comfy_io.Image.Output(),
-            ],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        model: str,
-        image: torch.Tensor,
-        prompt: str,
-        negative_prompt: str = "",
-        # width: int = 1024,
-        # height: int = 1024,
-        seed: int = 0,
-        watermark: bool = True,
-    ):
-        n_images = get_number_of_images(image)
-        if n_images not in (1, 2):
-            raise ValueError(f"Expected 1 or 2 input images, got {n_images}.")
-        images = []
-        for i in image:
-            images.append("data:image/png;base64," + tensor_to_base64_string(i, total_pixels=4096*4096))
-        payload = Image2ImageTaskCreationRequest(
-            model=model,
-            input=Image2ImageInputField(prompt=prompt, negative_prompt=negative_prompt, images=images),
-            parameters=Image2ImageParametersField(
-                # size=f"{width}*{height}",
-                seed=seed,
-                watermark=watermark,
-            ),
-        )
-        response = await process_task(
-            {
-                "auth_token": cls.hidden.auth_token_comfy_org,
-                "comfy_api_key": cls.hidden.api_key_comfy_org,
-            },
-            "/proxy/wan/api/v1/services/aigc/image2image/image-synthesis",
-            request_model=Image2ImageTaskCreationRequest,
-            response_model=ImageTaskStatusResponse,
-            payload=payload,
-            node_id=cls.hidden.unique_id,
-            estimated_duration=42,
-            poll_interval=3,
-        )
-        return comfy_io.NodeOutput(await download_url_to_image_tensor(str(response.output.results[0].url)))
-
-
 class WanTextToVideoApi(comfy_io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -739,7 +593,6 @@ class WanApiExtension(ComfyExtension):
     async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
         return [
             WanTextToImageApi,
-            WanImageToImageApi,
             WanTextToVideoApi,
             WanImageToVideoApi,
         ]

comfy_extras/nodes_audio.py

@@ -11,7 +11,6 @@ import json
 import random
 import hashlib
 import node_helpers
-import logging
 from comfy.cli_args import args
 from comfy.comfy_types import FileLocator
 
@@ -360,221 +359,11 @@ class RecordAudio:
     def load(self, audio):
         audio_path = folder_paths.get_annotated_filepath(audio)
 
-        waveform, sample_rate = load(audio_path)
+        waveform, sample_rate = torchaudio.load(audio_path)
         audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
         return (audio, )
 
 
-class TrimAudioDuration:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "audio": ("AUDIO",),
-                "start_index": ("FLOAT", {"default": 0.0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 0.01, "tooltip": "Start time in seconds, can be negative to count from the end (supports sub-seconds)."}),
-                "duration": ("FLOAT", {"default": 60.0, "min": 0.0, "step": 0.01, "tooltip": "Duration in seconds"}),
-            },
-        }
-
-    FUNCTION = "trim"
-    RETURN_TYPES = ("AUDIO",)
-    CATEGORY = "audio"
-    DESCRIPTION = "Trim audio tensor into chosen time range."
-
-    def trim(self, audio, start_index, duration):
-        waveform = audio["waveform"]
-        sample_rate = audio["sample_rate"]
-        audio_length = waveform.shape[-1]
-
-        if start_index < 0:
-            start_frame = audio_length + int(round(start_index * sample_rate))
-        else:
-            start_frame = int(round(start_index * sample_rate))
-        start_frame = max(0, min(start_frame, audio_length - 1))
-
-        end_frame = start_frame + int(round(duration * sample_rate))
-        end_frame = max(0, min(end_frame, audio_length))
-
-        if start_frame >= end_frame:
-            raise ValueError("AudioTrim: Start time must be less than end time and be within the audio length.")
-
-        return ({"waveform": waveform[..., start_frame:end_frame], "sample_rate": sample_rate},)
-
-
-class SplitAudioChannels:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "audio": ("AUDIO",),
-        }}
-
-    RETURN_TYPES = ("AUDIO", "AUDIO")
-    RETURN_NAMES = ("left", "right")
-    FUNCTION = "separate"
-    CATEGORY = "audio"
-    DESCRIPTION = "Separates the audio into left and right channels."
-
-    def separate(self, audio):
-        waveform = audio["waveform"]
-        sample_rate = audio["sample_rate"]
-
-        if waveform.shape[1] != 2:
-            raise ValueError("AudioSplit: Input audio has only one channel.")
-
-        left_channel = waveform[..., 0:1, :]
-        right_channel = waveform[..., 1:2, :]
-
-        return ({"waveform": left_channel, "sample_rate": sample_rate}, {"waveform": right_channel, "sample_rate": sample_rate})
-
-
-def match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2):
-    if sample_rate_1 != sample_rate_2:
-        if sample_rate_1 > sample_rate_2:
-            waveform_2 = torchaudio.functional.resample(waveform_2, sample_rate_2, sample_rate_1)
-            output_sample_rate = sample_rate_1
-            logging.info(f"Resampling audio2 from {sample_rate_2}Hz to {sample_rate_1}Hz for merging.")
-        else:
-            waveform_1 = torchaudio.functional.resample(waveform_1, sample_rate_1, sample_rate_2)
-            output_sample_rate = sample_rate_2
-            logging.info(f"Resampling audio1 from {sample_rate_1}Hz to {sample_rate_2}Hz for merging.")
-    else:
-        output_sample_rate = sample_rate_1
-    return waveform_1, waveform_2, output_sample_rate
-
-
-class AudioConcat:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "audio1": ("AUDIO",),
-            "audio2": ("AUDIO",),
-            "direction": (['after', 'before'], {"default": 'after', "tooltip": "Whether to append audio2 after or before audio1."}),
-        }}
-
-    RETURN_TYPES = ("AUDIO",)
-    FUNCTION = "concat"
-    CATEGORY = "audio"
-    DESCRIPTION = "Concatenates the audio1 to audio2 in the specified direction."
-
-    def concat(self, audio1, audio2, direction):
-        waveform_1 = audio1["waveform"]
-        waveform_2 = audio2["waveform"]
-        sample_rate_1 = audio1["sample_rate"]
-        sample_rate_2 = audio2["sample_rate"]
-
-        if waveform_1.shape[1] == 1:
-            waveform_1 = waveform_1.repeat(1, 2, 1)
-            logging.info("AudioConcat: Converted mono audio1 to stereo by duplicating the channel.")
-        if waveform_2.shape[1] == 1:
-            waveform_2 = waveform_2.repeat(1, 2, 1)
-            logging.info("AudioConcat: Converted mono audio2 to stereo by duplicating the channel.")
-
-        waveform_1, waveform_2, output_sample_rate = match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2)
-
-        if direction == 'after':
-            concatenated_audio = torch.cat((waveform_1, waveform_2), dim=2)
-        elif direction == 'before':
-            concatenated_audio = torch.cat((waveform_2, waveform_1), dim=2)
-
-        return ({"waveform": concatenated_audio, "sample_rate": output_sample_rate},)
-
-
-class AudioMerge:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "audio1": ("AUDIO",),
-                "audio2": ("AUDIO",),
-                "merge_method": (["add", "mean", "subtract", "multiply"], {"tooltip": "The method used to combine the audio waveforms."}),
-            },
-        }
-
-    FUNCTION = "merge"
-    RETURN_TYPES = ("AUDIO",)
-    CATEGORY = "audio"
-    DESCRIPTION = "Combine two audio tracks by overlaying their waveforms."
-
-    def merge(self, audio1, audio2, merge_method):
-        waveform_1 = audio1["waveform"]
-        waveform_2 = audio2["waveform"]
-        sample_rate_1 = audio1["sample_rate"]
-        sample_rate_2 = audio2["sample_rate"]
-
-        waveform_1, waveform_2, output_sample_rate = match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2)
-
-        length_1 = waveform_1.shape[-1]
-        length_2 = waveform_2.shape[-1]
-
-        if length_2 > length_1:
-            logging.info(f"AudioMerge: Trimming audio2 from {length_2} to {length_1} samples to match audio1 length.")
-            waveform_2 = waveform_2[..., :length_1]
-        elif length_2 < length_1:
-            logging.info(f"AudioMerge: Padding audio2 from {length_2} to {length_1} samples to match audio1 length.")
-            pad_shape = list(waveform_2.shape)
-            pad_shape[-1] = length_1 - length_2
-            pad_tensor = torch.zeros(pad_shape, dtype=waveform_2.dtype, device=waveform_2.device)
-            waveform_2 = torch.cat((waveform_2, pad_tensor), dim=-1)
-
-        if merge_method == "add":
-            waveform = waveform_1 + waveform_2
-        elif merge_method == "subtract":
-            waveform = waveform_1 - waveform_2
-        elif merge_method == "multiply":
-            waveform = waveform_1 * waveform_2
-        elif merge_method == "mean":
-            waveform = (waveform_1 + waveform_2) / 2
-
-        max_val = waveform.abs().max()
-        if max_val > 1.0:
-            waveform = waveform / max_val
-
-        return ({"waveform": waveform, "sample_rate": output_sample_rate},)
-
-
-class AudioAdjustVolume:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "audio": ("AUDIO",),
-            "volume": ("INT", {"default": 1.0, "min": -100, "max": 100, "tooltip": "Volume adjustment in decibels (dB). 0 = no change, +6 = double, -6 = half, etc"}),
-        }}
-
-    RETURN_TYPES = ("AUDIO",)
-    FUNCTION = "adjust_volume"
-    CATEGORY = "audio"
-
-    def adjust_volume(self, audio, volume):
-        if volume == 0:
-            return (audio,)
-        waveform = audio["waveform"]
-        sample_rate = audio["sample_rate"]
-
-        gain = 10 ** (volume / 20)
-        waveform = waveform * gain
-
-        return ({"waveform": waveform, "sample_rate": sample_rate},)
-
-
-class EmptyAudio:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "duration": ("FLOAT", {"default": 60.0, "min": 0.0, "max": 0xffffffffffffffff, "step": 0.01, "tooltip": "Duration of the empty audio clip in seconds"}),
-            "sample_rate": ("INT", {"default": 44100, "tooltip": "Sample rate of the empty audio clip."}),
-            "channels": ("INT", {"default": 2, "min": 1, "max": 2, "tooltip": "Number of audio channels (1 for mono, 2 for stereo)."}),
-        }}
-
-    RETURN_TYPES = ("AUDIO",)
-    FUNCTION = "create_empty_audio"
-    CATEGORY = "audio"
-
-    def create_empty_audio(self, duration, sample_rate, channels):
-        num_samples = int(round(duration * sample_rate))
-        waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
-        return ({"waveform": waveform, "sample_rate": sample_rate},)
-
-
 NODE_CLASS_MAPPINGS = {
     "EmptyLatentAudio": EmptyLatentAudio,
     "VAEEncodeAudio": VAEEncodeAudio,
@@ -586,12 +375,6 @@ NODE_CLASS_MAPPINGS = {
     "PreviewAudio": PreviewAudio,
     "ConditioningStableAudio": ConditioningStableAudio,
     "RecordAudio": RecordAudio,
-    "TrimAudioDuration": TrimAudioDuration,
-    "SplitAudioChannels": SplitAudioChannels,
-    "AudioConcat": AudioConcat,
-    "AudioMerge": AudioMerge,
-    "AudioAdjustVolume": AudioAdjustVolume,
-    "EmptyAudio": EmptyAudio,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -604,10 +387,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SaveAudioMP3": "Save Audio (MP3)",
     "SaveAudioOpus": "Save Audio (Opus)",
     "RecordAudio": "Record Audio",
-    "TrimAudioDuration": "Trim Audio Duration",
-    "SplitAudioChannels": "Split Audio Channels",
-    "AudioConcat": "Audio Concat",
-    "AudioMerge": "Audio Merge",
-    "AudioAdjustVolume": "Audio Adjust Volume",
-    "EmptyAudio": "Empty Audio",
 }

comfy_extras/nodes_audio_encoder.py

@@ -1,62 +1,44 @@
 import folder_paths
 import comfy.audio_encoders.audio_encoders
 import comfy.utils
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
-class AudioEncoderLoader(io.ComfyNode):
+class AudioEncoderLoader:
     @classmethod
-    def define_schema(cls) -> io.Schema:
-        return io.Schema(
-            node_id="AudioEncoderLoader",
-            category="loaders",
-            inputs=[
-                io.Combo.Input(
-                    "audio_encoder_name",
-                    options=folder_paths.get_filename_list("audio_encoders"),
-                ),
-            ],
-            outputs=[io.AudioEncoder.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "audio_encoder_name": (folder_paths.get_filename_list("audio_encoders"), ),
+                             }}
+    RETURN_TYPES = ("AUDIO_ENCODER",)
+    FUNCTION = "load_model"
 
-    @classmethod
-    def execute(cls, audio_encoder_name) -> io.NodeOutput:
+    CATEGORY = "loaders"
+
+    def load_model(self, audio_encoder_name):
         audio_encoder_name = folder_paths.get_full_path_or_raise("audio_encoders", audio_encoder_name)
         sd = comfy.utils.load_torch_file(audio_encoder_name, safe_load=True)
         audio_encoder = comfy.audio_encoders.audio_encoders.load_audio_encoder_from_sd(sd)
         if audio_encoder is None:
             raise RuntimeError("ERROR: audio encoder file is invalid and does not contain a valid model.")
-        return io.NodeOutput(audio_encoder)
+        return (audio_encoder,)
 
 
-class AudioEncoderEncode(io.ComfyNode):
+class AudioEncoderEncode:
     @classmethod
-    def define_schema(cls) -> io.Schema:
-        return io.Schema(
-            node_id="AudioEncoderEncode",
-            category="conditioning",
-            inputs=[
-                io.AudioEncoder.Input("audio_encoder"),
-                io.Audio.Input("audio"),
-            ],
-            outputs=[io.AudioEncoderOutput.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "audio_encoder": ("AUDIO_ENCODER",),
+                              "audio": ("AUDIO",),
+                             }}
+    RETURN_TYPES = ("AUDIO_ENCODER_OUTPUT",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, audio_encoder, audio) -> io.NodeOutput:
+    CATEGORY = "conditioning"
+
+    def encode(self, audio_encoder, audio):
         output = audio_encoder.encode_audio(audio["waveform"], audio["sample_rate"])
-        return io.NodeOutput(output)
+        return (output,)
 
 
-class AudioEncoder(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            AudioEncoderLoader,
-            AudioEncoderEncode,
-        ]
-
-
-async def comfy_entrypoint() -> AudioEncoder:
-    return AudioEncoder()
+NODE_CLASS_MAPPINGS = {
+    "AudioEncoderLoader": AudioEncoderLoader,
+    "AudioEncoderEncode": AudioEncoderEncode,
+}

comfy_extras/nodes_clip_sdxl.py

@@ -1,52 +1,43 @@
-from typing_extensions import override
+from nodes import MAX_RESOLUTION
 
-import nodes
-from comfy_api.latest import ComfyExtension, io
-
-
-class CLIPTextEncodeSDXLRefiner(io.ComfyNode):
+class CLIPTextEncodeSDXLRefiner:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeSDXLRefiner",
-            category="advanced/conditioning",
-            inputs=[
-                io.Float.Input("ascore", default=6.0, min=0.0, max=1000.0, step=0.01),
-                io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.String.Input("text", multiline=True, dynamic_prompts=True),
-                io.Clip.Input("clip"),
-            ],
-            outputs=[io.Conditioning.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "ascore": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
+            "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ),
+            }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, clip, ascore, width, height, text) -> io.NodeOutput:
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, ascore, width, height, text):
         tokens = clip.tokenize(text)
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height}))
+        return (clip.encode_from_tokens_scheduled(tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height}), )
 
-class CLIPTextEncodeSDXL(io.ComfyNode):
+class CLIPTextEncodeSDXL:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeSDXL",
-            category="advanced/conditioning",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("crop_w", default=0, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("crop_h", default=0, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("target_width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("target_height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.String.Input("text_g", multiline=True, dynamic_prompts=True),
-                io.String.Input("text_l", multiline=True, dynamic_prompts=True),
-            ],
-            outputs=[io.Conditioning.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+            "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+            "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "text_g": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "text_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l) -> io.NodeOutput:
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l):
         tokens = clip.tokenize(text_g)
         tokens["l"] = clip.tokenize(text_l)["l"]
         if len(tokens["l"]) != len(tokens["g"]):
@@ -55,17 +46,9 @@ class CLIPTextEncodeSDXL(io.ComfyNode):
                 tokens["l"] += empty["l"]
             while len(tokens["l"]) > len(tokens["g"]):
                 tokens["g"] += empty["g"]
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}))
+        return (clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}), )
 
-
-class ClipSdxlExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            CLIPTextEncodeSDXLRefiner,
-            CLIPTextEncodeSDXL,
-        ]
-
-
-async def comfy_entrypoint() -> ClipSdxlExtension:
-    return ClipSdxlExtension()
+NODE_CLASS_MAPPINGS = {
+    "CLIPTextEncodeSDXLRefiner": CLIPTextEncodeSDXLRefiner,
+    "CLIPTextEncodeSDXL": CLIPTextEncodeSDXL,
+}

comfy_extras/nodes_compositing.py

@@ -1,9 +1,6 @@
 import torch
 import comfy.utils
 from enum import Enum
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
-
 
 def resize_mask(mask, shape):
     return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
@@ -104,28 +101,24 @@ def porter_duff_composite(src_image: torch.Tensor, src_alpha: torch.Tensor, dst_
     return out_image, out_alpha
 
 
-class PorterDuffImageComposite(io.ComfyNode):
+class PorterDuffImageComposite:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PorterDuffImageComposite",
-            display_name="Porter-Duff Image Composite",
-            category="mask/compositing",
-            inputs=[
-                io.Image.Input("source"),
-                io.Mask.Input("source_alpha"),
-                io.Image.Input("destination"),
-                io.Mask.Input("destination_alpha"),
-                io.Combo.Input("mode", options=[mode.name for mode in PorterDuffMode], default=PorterDuffMode.DST.name),
-            ],
-            outputs=[
-                io.Image.Output(),
-                io.Mask.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "source": ("IMAGE",),
+                "source_alpha": ("MASK",),
+                "destination": ("IMAGE",),
+                "destination_alpha": ("MASK",),
+                "mode": ([mode.name for mode in PorterDuffMode], {"default": PorterDuffMode.DST.name}),
+            },
+        }
 
-    @classmethod
-    def execute(cls, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode) -> io.NodeOutput:
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "composite"
+    CATEGORY = "mask/compositing"
+
+    def composite(self, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode):
         batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha))
         out_images = []
         out_alphas = []
@@ -157,48 +150,45 @@ class PorterDuffImageComposite(io.ComfyNode):
             out_images.append(out_image)
             out_alphas.append(out_alpha.squeeze(2))
 
-        return io.NodeOutput(torch.stack(out_images), torch.stack(out_alphas))
+        result = (torch.stack(out_images), torch.stack(out_alphas))
+        return result
 
 
-class SplitImageWithAlpha(io.ComfyNode):
+class SplitImageWithAlpha:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SplitImageWithAlpha",
-            display_name="Split Image with Alpha",
-            category="mask/compositing",
-            inputs=[
-                io.Image.Input("image"),
-            ],
-            outputs=[
-                io.Image.Output(),
-                io.Mask.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                    "image": ("IMAGE",),
+                }
+        }
 
-    @classmethod
-    def execute(cls, image: torch.Tensor) -> io.NodeOutput:
+    CATEGORY = "mask/compositing"
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "split_image_with_alpha"
+
+    def split_image_with_alpha(self, image: torch.Tensor):
         out_images = [i[:,:,:3] for i in image]
         out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image]
-        return io.NodeOutput(torch.stack(out_images), 1.0 - torch.stack(out_alphas))
+        result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas))
+        return result
 
 
-class JoinImageWithAlpha(io.ComfyNode):
+class JoinImageWithAlpha:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="JoinImageWithAlpha",
-            display_name="Join Image with Alpha",
-            category="mask/compositing",
-            inputs=[
-                io.Image.Input("image"),
-                io.Mask.Input("alpha"),
-            ],
-            outputs=[io.Image.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                    "image": ("IMAGE",),
+                    "alpha": ("MASK",),
+                }
+        }
 
-    @classmethod
-    def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput:
+    CATEGORY = "mask/compositing"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "join_image_with_alpha"
+
+    def join_image_with_alpha(self, image: torch.Tensor, alpha: torch.Tensor):
         batch_size = min(len(image), len(alpha))
         out_images = []
 
@@ -206,18 +196,19 @@ class JoinImageWithAlpha(io.ComfyNode):
         for i in range(batch_size):
            out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
 
-        return io.NodeOutput(torch.stack(out_images))
+        result = (torch.stack(out_images),)
+        return result
 
 
-class CompositingExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            PorterDuffImageComposite,
-            SplitImageWithAlpha,
-            JoinImageWithAlpha,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "PorterDuffImageComposite": PorterDuffImageComposite,
+    "SplitImageWithAlpha": SplitImageWithAlpha,
+    "JoinImageWithAlpha": JoinImageWithAlpha,
+}
 
 
-async def comfy_entrypoint() -> CompositingExtension:
-    return CompositingExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PorterDuffImageComposite": "Porter-Duff Image Composite",
+    "SplitImageWithAlpha": "Split Image with Alpha",
+    "JoinImageWithAlpha": "Join Image with Alpha",
+}

comfy_extras/nodes_differential_diffusion.py

@@ -1,41 +1,34 @@
 # code adapted from https://github.com/exx8/differential-diffusion
 
-from typing_extensions import override
-
 import torch
-from comfy_api.latest import ComfyExtension, io
 
-
-class DifferentialDiffusion(io.ComfyNode):
+class DifferentialDiffusion():
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="DifferentialDiffusion",
-            display_name="Differential Diffusion",
-            category="_for_testing",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input(
-                    "strength",
-                    default=1.0,
-                    min=0.0,
-                    max=1.0,
-                    step=0.01,
-                    optional=True,
-                ),
-            ],
-            outputs=[io.Model.Output()],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+            },
+            "optional": {
+                "strength": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                }),
+            }
+        }
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply"
+    CATEGORY = "_for_testing"
+    INIT = False
 
-    @classmethod
-    def execute(cls, model, strength=1.0) -> io.NodeOutput:
+    def apply(self, model, strength=1.0):
         model = model.clone()
-        model.set_model_denoise_mask_function(lambda *args, **kwargs: cls.forward(*args, **kwargs, strength=strength))
-        return io.NodeOutput(model)
+        model.set_model_denoise_mask_function(lambda *args, **kwargs: self.forward(*args, **kwargs, strength=strength))
+        return (model, )
 
-    @classmethod
-    def forward(cls, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):
+    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):
         model = extra_options["model"]
         step_sigmas = extra_options["sigmas"]
         sigma_to = model.inner_model.model_sampling.sigma_min
@@ -60,13 +53,9 @@ class DifferentialDiffusion(io.ComfyNode):
             return binary_mask
 
 
-class DifferentialDiffusionExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            DifferentialDiffusion,
-        ]
-
-
-async def comfy_entrypoint() -> DifferentialDiffusionExtension:
-    return DifferentialDiffusionExtension()
+NODE_CLASS_MAPPINGS = {
+    "DifferentialDiffusion": DifferentialDiffusion,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DifferentialDiffusion": "Differential Diffusion",
+}

comfy_extras/nodes_edit_model.py

@@ -1,38 +1,26 @@
 import node_helpers
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
-class ReferenceLatent(io.ComfyNode):
+class ReferenceLatent:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ReferenceLatent",
-            category="advanced/conditioning/edit_models",
-            description="This node sets the guiding latent for an edit model. If the model supports it you can chain multiple to set multiple reference images.",
-            inputs=[
-                io.Conditioning.Input("conditioning"),
-                io.Latent.Input("latent", optional=True),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"conditioning": ("CONDITIONING", ),
+                            },
+                "optional": {"latent": ("LATENT", ),}
+               }
 
-    @classmethod
-    def execute(cls, conditioning, latent=None) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "append"
+
+    CATEGORY = "advanced/conditioning/edit_models"
+    DESCRIPTION = "This node sets the guiding latent for an edit model. If the model supports it you can chain multiple to set multiple reference images."
+
+    def append(self, conditioning, latent=None):
         if latent is not None:
             conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": [latent["samples"]]}, append=True)
-        return io.NodeOutput(conditioning)
+        return (conditioning, )
 
 
-class EditModelExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            ReferenceLatent,
-        ]
-
-
-def comfy_entrypoint() -> EditModelExtension:
-    return EditModelExtension()
+NODE_CLASS_MAPPINGS = {
+    "ReferenceLatent": ReferenceLatent,
+}

comfy_extras/nodes_eps.py

@@ -1,74 +0,0 @@
-from typing_extensions import override
-
-from comfy_api.latest import ComfyExtension, io
-
-
-class EpsilonScaling(io.ComfyNode):
-    """
-    Implements the Epsilon Scaling method from 'Elucidating the Exposure Bias in Diffusion Models'
-    (https://arxiv.org/abs/2308.15321v6).
-
-    This method mitigates exposure bias by scaling the predicted noise during sampling,
-    which can significantly improve sample quality. This implementation uses the "uniform schedule"
-    recommended by the paper for its practicality and effectiveness.
-    """
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="Epsilon Scaling",
-            category="model_patches/unet",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input(
-                    "scaling_factor",
-                    default=1.005,
-                    min=0.5,
-                    max=1.5,
-                    step=0.001,
-                    display_mode=io.NumberDisplay.number,
-                ),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
-
-    @classmethod
-    def execute(cls, model, scaling_factor) -> io.NodeOutput:
-        # Prevent division by zero, though the UI's min value should prevent this.
-        if scaling_factor == 0:
-            scaling_factor = 1e-9
-
-        def epsilon_scaling_function(args):
-            """
-            This function is applied after the CFG guidance has been calculated.
-            It recalculates the denoised latent by scaling the predicted noise.
-            """
-            denoised = args["denoised"]
-            x = args["input"]
-
-            noise_pred = x - denoised
-
-            scaled_noise_pred = noise_pred / scaling_factor
-
-            new_denoised = x - scaled_noise_pred
-
-            return new_denoised
-
-        # Clone the model patcher to avoid modifying the original model in place
-        model_clone = model.clone()
-
-        model_clone.set_model_sampler_post_cfg_function(epsilon_scaling_function)
-
-        return io.NodeOutput(model_clone)
-
-
-class EpsilonScalingExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            EpsilonScaling,
-        ]
-
-async def comfy_entrypoint() -> EpsilonScalingExtension:
-    return EpsilonScalingExtension()

comfy_extras/nodes_flux.py

@@ -1,80 +1,60 @@
 import node_helpers
 import comfy.utils
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
-
-class CLIPTextEncodeFlux(io.ComfyNode):
+class CLIPTextEncodeFlux:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeFlux",
-            category="advanced/conditioning/flux",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.String.Input("clip_l", multiline=True, dynamic_prompts=True),
-                io.String.Input("t5xxl", multiline=True, dynamic_prompts=True),
-                io.Float.Input("guidance", default=3.5, min=0.0, max=100.0, step=0.1),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "clip_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "t5xxl": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
+            }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, clip, clip_l, t5xxl, guidance) -> io.NodeOutput:
+    CATEGORY = "advanced/conditioning/flux"
+
+    def encode(self, clip, clip_l, t5xxl, guidance):
         tokens = clip.tokenize(clip_l)
         tokens["t5xxl"] = clip.tokenize(t5xxl)["t5xxl"]
 
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"guidance": guidance}))
+        return (clip.encode_from_tokens_scheduled(tokens, add_dict={"guidance": guidance}), )
 
-    encode = execute  # TODO: remove
-
-
-class FluxGuidance(io.ComfyNode):
+class FluxGuidance:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="FluxGuidance",
-            category="advanced/conditioning/flux",
-            inputs=[
-                io.Conditioning.Input("conditioning"),
-                io.Float.Input("guidance", default=3.5, min=0.0, max=100.0, step=0.1),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "conditioning": ("CONDITIONING", ),
+            "guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
+            }}
 
-    @classmethod
-    def execute(cls, conditioning, guidance) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "append"
+
+    CATEGORY = "advanced/conditioning/flux"
+
+    def append(self, conditioning, guidance):
         c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
-        return io.NodeOutput(c)
-
-    append = execute  # TODO: remove
+        return (c, )
 
 
-class FluxDisableGuidance(io.ComfyNode):
+class FluxDisableGuidance:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="FluxDisableGuidance",
-            category="advanced/conditioning/flux",
-            description="This node completely disables the guidance embed on Flux and Flux like models",
-            inputs=[
-                io.Conditioning.Input("conditioning"),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "conditioning": ("CONDITIONING", ),
+            }}
 
-    @classmethod
-    def execute(cls, conditioning) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "append"
+
+    CATEGORY = "advanced/conditioning/flux"
+    DESCRIPTION = "This node completely disables the guidance embed on Flux and Flux like models"
+
+    def append(self, conditioning):
         c = node_helpers.conditioning_set_values(conditioning, {"guidance": None})
-        return io.NodeOutput(c)
-
-    append = execute  # TODO: remove
+        return (c, )
 
 
 PREFERED_KONTEXT_RESOLUTIONS = [
@@ -98,73 +78,52 @@ PREFERED_KONTEXT_RESOLUTIONS = [
 ]
 
 
-class FluxKontextImageScale(io.ComfyNode):
+class FluxKontextImageScale:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="FluxKontextImageScale",
-            category="advanced/conditioning/flux",
-            description="This node resizes the image to one that is more optimal for flux kontext.",
-            inputs=[
-                io.Image.Input("image"),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"image": ("IMAGE", ),
+                            },
+               }
 
-    @classmethod
-    def execute(cls, image) -> io.NodeOutput:
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "scale"
+
+    CATEGORY = "advanced/conditioning/flux"
+    DESCRIPTION = "This node resizes the image to one that is more optimal for flux kontext."
+
+    def scale(self, image):
         width = image.shape[2]
         height = image.shape[1]
         aspect_ratio = width / height
         _, width, height = min((abs(aspect_ratio - w / h), w, h) for w, h in PREFERED_KONTEXT_RESOLUTIONS)
         image = comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "lanczos", "center").movedim(1, -1)
-        return io.NodeOutput(image)
-
-    scale = execute  # TODO: remove
+        return (image, )
 
 
-class FluxKontextMultiReferenceLatentMethod(io.ComfyNode):
+class FluxKontextMultiReferenceLatentMethod:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="FluxKontextMultiReferenceLatentMethod",
-            category="advanced/conditioning/flux",
-            inputs=[
-                io.Conditioning.Input("conditioning"),
-                io.Combo.Input(
-                    "reference_latents_method",
-                    options=["offset", "index", "uxo/uno"],
-                ),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "conditioning": ("CONDITIONING", ),
+            "reference_latents_method": (("offset", "index", "uxo/uno"), ),
+            }}
 
-    @classmethod
-    def execute(cls, conditioning, reference_latents_method) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "append"
+    EXPERIMENTAL = True
+
+    CATEGORY = "advanced/conditioning/flux"
+
+    def append(self, conditioning, reference_latents_method):
         if "uxo" in reference_latents_method or "uso" in reference_latents_method:
             reference_latents_method = "uxo"
         c = node_helpers.conditioning_set_values(conditioning, {"reference_latents_method": reference_latents_method})
-        return io.NodeOutput(c)
+        return (c, )
 
-    append = execute  # TODO: remove
-
-
-class FluxExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            CLIPTextEncodeFlux,
-            FluxGuidance,
-            FluxDisableGuidance,
-            FluxKontextImageScale,
-            FluxKontextMultiReferenceLatentMethod,
-        ]
-
-
-async def comfy_entrypoint() -> FluxExtension:
-    return FluxExtension()
+NODE_CLASS_MAPPINGS = {
+    "CLIPTextEncodeFlux": CLIPTextEncodeFlux,
+    "FluxGuidance": FluxGuidance,
+    "FluxDisableGuidance": FluxDisableGuidance,
+    "FluxKontextImageScale": FluxKontextImageScale,
+    "FluxKontextMultiReferenceLatentMethod": FluxKontextMultiReferenceLatentMethod,
+}

comfy_extras/nodes_fresca.py

@@ -1,8 +1,6 @@
 # Code based on https://github.com/WikiChao/FreSca (MIT License)
 import torch
 import torch.fft as fft
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
 def Fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20):
@@ -53,31 +51,25 @@ def Fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20):
     return x_filtered
 
 
-class FreSca(io.ComfyNode):
+class FreSca:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="FreSca",
-            display_name="FreSca",
-            category="_for_testing",
-            description="Applies frequency-dependent scaling to the guidance",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("scale_low", default=1.0, min=0, max=10, step=0.01,
-                               tooltip="Scaling factor for low-frequency components"),
-                io.Float.Input("scale_high", default=1.25, min=0, max=10, step=0.01,
-                               tooltip="Scaling factor for high-frequency components"),
-                io.Int.Input("freq_cutoff", default=20, min=1, max=10000, step=1,
-                             tooltip="Number of frequency indices around center to consider as low-frequency"),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-            is_experimental=True,
-        )
-
-    @classmethod
-    def execute(cls, model, scale_low, scale_high, freq_cutoff):
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "scale_low": ("FLOAT", {"default": 1.0, "min": 0, "max": 10, "step": 0.01,
+                                        "tooltip": "Scaling factor for low-frequency components"}),
+                "scale_high": ("FLOAT", {"default": 1.25, "min": 0, "max": 10, "step": 0.01,
+                                        "tooltip": "Scaling factor for high-frequency components"}),
+                "freq_cutoff": ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1,
+                                        "tooltip": "Number of frequency indices around center to consider as low-frequency"}),
+            }
+        }
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+    CATEGORY = "_for_testing"
+    DESCRIPTION = "Applies frequency-dependent scaling to the guidance"
+    def patch(self, model, scale_low, scale_high, freq_cutoff):
         def custom_cfg_function(args):
             conds_out = args["conds_out"]
             if len(conds_out) <= 1 or None in args["conds"][:2]:
@@ -99,16 +91,13 @@ class FreSca(io.ComfyNode):
         m = model.clone()
         m.set_model_sampler_pre_cfg_function(custom_cfg_function)
 
-        return io.NodeOutput(m)
+        return (m,)
 
 
-class FreScaExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            FreSca,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "FreSca": FreSca,
+}
 
-
-async def comfy_entrypoint() -> FreScaExtension:
-    return FreScaExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FreSca": "FreSca",
+}

comfy_extras/nodes_gits.py

@@ -1,8 +1,6 @@
 # from https://github.com/zju-pi/diff-sampler/tree/main/gits-main
 import numpy as np
 import torch
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 def loglinear_interp(t_steps, num_steps):
     """
@@ -335,28 +333,25 @@ NOISE_LEVELS = {
     ],
 }
 
-class GITSScheduler(io.ComfyNode):
+class GITSScheduler:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="GITSScheduler",
-            category="sampling/custom_sampling/schedulers",
-            inputs=[
-                io.Float.Input("coeff", default=1.20, min=0.80, max=1.50, step=0.05),
-                io.Int.Input("steps", default=10, min=2, max=1000),
-                io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
-            ],
-            outputs=[
-                io.Sigmas.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"coeff": ("FLOAT", {"default": 1.20, "min": 0.80, "max": 1.50, "step": 0.05}),
+                     "steps": ("INT", {"default": 10, "min": 2, "max": 1000}),
+                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
 
-    @classmethod
-    def execute(cls, coeff, steps, denoise):
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, coeff, steps, denoise):
         total_steps = steps
         if denoise < 1.0:
             if denoise <= 0.0:
-                return io.NodeOutput(torch.FloatTensor([]))
+                return (torch.FloatTensor([]),)
             total_steps = round(steps * denoise)
 
         if steps <= 20:
@@ -367,16 +362,8 @@ class GITSScheduler(io.ComfyNode):
 
         sigmas = sigmas[-(total_steps + 1):]
         sigmas[-1] = 0
-        return io.NodeOutput(torch.FloatTensor(sigmas))
+        return (torch.FloatTensor(sigmas), )
 
-
-class GITSSchedulerExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            GITSScheduler,
-        ]
-
-
-async def comfy_entrypoint() -> GITSSchedulerExtension:
-    return GITSSchedulerExtension()
+NODE_CLASS_MAPPINGS = {
+    "GITSScheduler": GITSScheduler,
+}

comfy_extras/nodes_hidream.py

@@ -1,73 +1,55 @@
-from typing_extensions import override
-
 import folder_paths
 import comfy.sd
 import comfy.model_management
-from comfy_api.latest import ComfyExtension, io
 
 
-class QuadrupleCLIPLoader(io.ComfyNode):
+class QuadrupleCLIPLoader:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="QuadrupleCLIPLoader",
-            category="advanced/loaders",
-            description="[Recipes]\n\nhidream: long clip-l, long clip-g, t5xxl, llama_8b_3.1_instruct",
-            inputs=[
-                io.Combo.Input("clip_name1", options=folder_paths.get_filename_list("text_encoders")),
-                io.Combo.Input("clip_name2", options=folder_paths.get_filename_list("text_encoders")),
-                io.Combo.Input("clip_name3", options=folder_paths.get_filename_list("text_encoders")),
-                io.Combo.Input("clip_name4", options=folder_paths.get_filename_list("text_encoders")),
-            ],
-            outputs=[
-                io.Clip.Output(),
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "clip_name1": (folder_paths.get_filename_list("text_encoders"), ),
+                              "clip_name2": (folder_paths.get_filename_list("text_encoders"), ),
+                              "clip_name3": (folder_paths.get_filename_list("text_encoders"), ),
+                              "clip_name4": (folder_paths.get_filename_list("text_encoders"), )
+                             }}
+    RETURN_TYPES = ("CLIP",)
+    FUNCTION = "load_clip"
 
-    @classmethod
-    def execute(cls, clip_name1, clip_name2, clip_name3, clip_name4):
+    CATEGORY = "advanced/loaders"
+
+    DESCRIPTION = "[Recipes]\n\nhidream: long clip-l, long clip-g, t5xxl, llama_8b_3.1_instruct"
+
+    def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4):
         clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
         clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
         clip_path3 = folder_paths.get_full_path_or_raise("text_encoders", clip_name3)
         clip_path4 = folder_paths.get_full_path_or_raise("text_encoders", clip_name4)
         clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2, clip_path3, clip_path4], embedding_directory=folder_paths.get_folder_paths("embeddings"))
-        return io.NodeOutput(clip)
+        return (clip,)
 
-class CLIPTextEncodeHiDream(io.ComfyNode):
+class CLIPTextEncodeHiDream:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeHiDream",
-            category="advanced/conditioning",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.String.Input("clip_l", multiline=True, dynamic_prompts=True),
-                io.String.Input("clip_g", multiline=True, dynamic_prompts=True),
-                io.String.Input("t5xxl", multiline=True, dynamic_prompts=True),
-                io.String.Input("llama", multiline=True, dynamic_prompts=True),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "clip_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "clip_g": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "t5xxl": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "llama": ("STRING", {"multiline": True, "dynamicPrompts": True})
+            }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, clip_l, clip_g, t5xxl, llama):
 
-    @classmethod
-    def execute(cls, clip, clip_l, clip_g, t5xxl, llama):
         tokens = clip.tokenize(clip_g)
         tokens["l"] = clip.tokenize(clip_l)["l"]
         tokens["t5xxl"] = clip.tokenize(t5xxl)["t5xxl"]
         tokens["llama"] = clip.tokenize(llama)["llama"]
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens))
+        return (clip.encode_from_tokens_scheduled(tokens), )
 
-
-class HiDreamExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            QuadrupleCLIPLoader,
-            CLIPTextEncodeHiDream,
-        ]
-
-
-async def comfy_entrypoint() -> HiDreamExtension:
-    return HiDreamExtension()
+NODE_CLASS_MAPPINGS = {
+    "QuadrupleCLIPLoader": QuadrupleCLIPLoader,
+    "CLIPTextEncodeHiDream": CLIPTextEncodeHiDream,
+}

comfy_extras/nodes_hypertile.py

@@ -1,11 +1,9 @@
 #Taken from: https://github.com/tfernd/HyperTile/
 
 import math
-from typing_extensions import override
 from einops import rearrange
 # Use torch rng for consistency across generations
 from torch import randint
-from comfy_api.latest import ComfyExtension, io
 
 def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
     min_value = min(min_value, value)
@@ -22,31 +20,25 @@ def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
 
     return ns[idx]
 
-class HyperTile(io.ComfyNode):
+class HyperTile:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="HyperTile",
-            category="model_patches/unet",
-            inputs=[
-                io.Model.Input("model"),
-                io.Int.Input("tile_size", default=256, min=1, max=2048),
-                io.Int.Input("swap_size", default=2, min=1, max=128),
-                io.Int.Input("max_depth", default=0, min=0, max=10),
-                io.Boolean.Input("scale_depth", default=False),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                             "tile_size": ("INT", {"default": 256, "min": 1, "max": 2048}),
+                             "swap_size": ("INT", {"default": 2, "min": 1, "max": 128}),
+                             "max_depth": ("INT", {"default": 0, "min": 0, "max": 10}),
+                             "scale_depth": ("BOOLEAN", {"default": False}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, tile_size, swap_size, max_depth, scale_depth) -> io.NodeOutput:
+    CATEGORY = "model_patches/unet"
+
+    def patch(self, model, tile_size, swap_size, max_depth, scale_depth):
         latent_tile_size = max(32, tile_size) // 8
-        temp = None
+        self.temp = None
 
         def hypertile_in(q, k, v, extra_options):
-            nonlocal temp
             model_chans = q.shape[-2]
             orig_shape = extra_options['original_shape']
             apply_to = []
@@ -66,15 +58,14 @@ class HyperTile(io.ComfyNode):
 
                 if nh * nw > 1:
                     q = rearrange(q, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw)
-                    temp = (nh, nw, h, w)
+                    self.temp = (nh, nw, h, w)
                 return q, k, v
 
             return q, k, v
         def hypertile_out(out, extra_options):
-            nonlocal temp
-            if temp is not None:
-                nh, nw, h, w = temp
-                temp = None
+            if self.temp is not None:
+                nh, nw, h, w = self.temp
+                self.temp = None
                 out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw)
                 out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw)
             return out
@@ -85,14 +76,6 @@ class HyperTile(io.ComfyNode):
         m.set_model_attn1_output_patch(hypertile_out)
         return (m, )
 
-
-class HyperTileExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            HyperTile,
-        ]
-
-
-async def comfy_entrypoint() -> HyperTileExtension:
-    return HyperTileExtension()
+NODE_CLASS_MAPPINGS = {
+    "HyperTile": HyperTile,
+}

comfy_extras/nodes_ip2p.py

@@ -1,30 +1,21 @@
 import torch
 
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
-
-
-class InstructPixToPixConditioning(io.ComfyNode):
+class InstructPixToPixConditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="InstructPixToPixConditioning",
-            category="conditioning/instructpix2pix",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Vae.Input("vae"),
-                io.Image.Input("pixels"),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "vae": ("VAE", ),
+                             "pixels": ("IMAGE", ),
+                             }}
 
-    @classmethod
-    def execute(cls, positive, negative, pixels, vae) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/instructpix2pix"
+
+    def encode(self, positive, negative, pixels, vae):
         x = (pixels.shape[1] // 8) * 8
         y = (pixels.shape[2] // 8) * 8
 
@@ -47,17 +38,8 @@ class InstructPixToPixConditioning(io.ComfyNode):
                 n = [t[0], d]
                 c.append(n)
             out.append(c)
-        return io.NodeOutput(out[0], out[1], out_latent)
-
-
-class InstructPix2PixExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            InstructPixToPixConditioning,
-        ]
-
-
-async def comfy_entrypoint() -> InstructPix2PixExtension:
-    return InstructPix2PixExtension()
+        return (out[0], out[1], out_latent)
 
+NODE_CLASS_MAPPINGS = {
+    "InstructPixToPixConditioning": InstructPixToPixConditioning,
+}

comfy_extras/nodes_latent.py

@@ -2,8 +2,6 @@ import comfy.utils
 import comfy_extras.nodes_post_processing
 import torch
 import nodes
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
 def reshape_latent_to(target_shape, latent, repeat_batch=True):
@@ -15,23 +13,17 @@ def reshape_latent_to(target_shape, latent, repeat_batch=True):
         return latent
 
 
-class LatentAdd(io.ComfyNode):
+class LatentAdd:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentAdd",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples1"),
-                io.Latent.Input("samples2"),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
 
-    @classmethod
-    def execute(cls, samples1, samples2) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples1, samples2):
         samples_out = samples1.copy()
 
         s1 = samples1["samples"]
@@ -39,25 +31,19 @@ class LatentAdd(io.ComfyNode):
 
         s2 = reshape_latent_to(s1.shape, s2)
         samples_out["samples"] = s1 + s2
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentSubtract(io.ComfyNode):
+class LatentSubtract:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentSubtract",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples1"),
-                io.Latent.Input("samples2"),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
 
-    @classmethod
-    def execute(cls, samples1, samples2) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples1, samples2):
         samples_out = samples1.copy()
 
         s1 = samples1["samples"]
@@ -65,49 +51,41 @@ class LatentSubtract(io.ComfyNode):
 
         s2 = reshape_latent_to(s1.shape, s2)
         samples_out["samples"] = s1 - s2
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentMultiply(io.ComfyNode):
+class LatentMultiply:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentMultiply",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples"),
-                io.Float.Input("multiplier", default=1.0, min=-10.0, max=10.0, step=0.01),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples": ("LATENT",),
+                              "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                             }}
 
-    @classmethod
-    def execute(cls, samples, multiplier) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples, multiplier):
         samples_out = samples.copy()
 
         s1 = samples["samples"]
         samples_out["samples"] = s1 * multiplier
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentInterpolate(io.ComfyNode):
+class LatentInterpolate:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentInterpolate",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples1"),
-                io.Latent.Input("samples2"),
-                io.Float.Input("ratio", default=1.0, min=0.0, max=1.0, step=0.01),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples1": ("LATENT",),
+                              "samples2": ("LATENT",),
+                              "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                              }}
 
-    @classmethod
-    def execute(cls, samples1, samples2, ratio) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples1, samples2, ratio):
         samples_out = samples1.copy()
 
         s1 = samples1["samples"]
@@ -126,26 +104,19 @@ class LatentInterpolate(io.ComfyNode):
         st = torch.nan_to_num(t / mt)
 
         samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio))
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentConcat(io.ComfyNode):
+class LatentConcat:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentConcat",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples1"),
-                io.Latent.Input("samples2"),
-                io.Combo.Input("dim", options=["x", "-x", "y", "-y", "t", "-t"]),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",), "dim": (["x", "-x", "y", "-y", "t", "-t"], )}}
 
-    @classmethod
-    def execute(cls, samples1, samples2, dim) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples1, samples2, dim):
         samples_out = samples1.copy()
 
         s1 = samples1["samples"]
@@ -165,27 +136,22 @@ class LatentConcat(io.ComfyNode):
             dim = -3
 
         samples_out["samples"] = torch.cat(c, dim=dim)
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentCut(io.ComfyNode):
+class LatentCut:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentCut",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples"),
-                io.Combo.Input("dim", options=["x", "y", "t"]),
-                io.Int.Input("index", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, step=1),
-                io.Int.Input("amount", default=1, min=1, max=nodes.MAX_RESOLUTION, step=1),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"samples": ("LATENT",),
+                             "dim": (["x", "y", "t"], ),
+                             "index": ("INT", {"default": 0, "min": -nodes.MAX_RESOLUTION, "max": nodes.MAX_RESOLUTION, "step": 1}),
+                             "amount": ("INT", {"default": 1, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 1})}}
 
-    @classmethod
-    def execute(cls, samples, dim, index, amount) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples, dim, index, amount):
         samples_out = samples.copy()
 
         s1 = samples["samples"]
@@ -205,25 +171,19 @@ class LatentCut(io.ComfyNode):
             amount = min(-index, amount)
 
         samples_out["samples"] = torch.narrow(s1, dim, index, amount)
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentBatch(io.ComfyNode):
+class LatentBatch:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentBatch",
-            category="latent/batch",
-            inputs=[
-                io.Latent.Input("samples1"),
-                io.Latent.Input("samples2"),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
 
-    @classmethod
-    def execute(cls, samples1, samples2) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "batch"
+
+    CATEGORY = "latent/batch"
+
+    def batch(self, samples1, samples2):
         samples_out = samples1.copy()
         s1 = samples1["samples"]
         s2 = samples2["samples"]
@@ -232,25 +192,20 @@ class LatentBatch(io.ComfyNode):
         s = torch.cat((s1, s2), dim=0)
         samples_out["samples"] = s
         samples_out["batch_index"] = samples1.get("batch_index", [x for x in range(0, s1.shape[0])]) + samples2.get("batch_index", [x for x in range(0, s2.shape[0])])
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentBatchSeedBehavior(io.ComfyNode):
+class LatentBatchSeedBehavior:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentBatchSeedBehavior",
-            category="latent/advanced",
-            inputs=[
-                io.Latent.Input("samples"),
-                io.Combo.Input("seed_behavior", options=["random", "fixed"], default="fixed"),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples": ("LATENT",),
+                              "seed_behavior": (["random", "fixed"],{"default": "fixed"}),}}
 
-    @classmethod
-    def execute(cls, samples, seed_behavior) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced"
+
+    def op(self, samples, seed_behavior):
         samples_out = samples.copy()
         latent = samples["samples"]
         if seed_behavior == "random":
@@ -260,50 +215,41 @@ class LatentBatchSeedBehavior(io.ComfyNode):
             batch_number = samples_out.get("batch_index", [0])[0]
             samples_out["batch_index"] = [batch_number] * latent.shape[0]
 
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentApplyOperation(io.ComfyNode):
+class LatentApplyOperation:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentApplyOperation",
-            category="latent/advanced/operations",
-            is_experimental=True,
-            inputs=[
-                io.Latent.Input("samples"),
-                io.LatentOperation.Input("operation"),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "samples": ("LATENT",),
+                             "operation": ("LATENT_OPERATION",),
+                             }}
 
-    @classmethod
-    def execute(cls, samples, operation) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced/operations"
+    EXPERIMENTAL = True
+
+    def op(self, samples, operation):
         samples_out = samples.copy()
 
         s1 = samples["samples"]
         samples_out["samples"] = operation(latent=s1)
-        return io.NodeOutput(samples_out)
+        return (samples_out,)
 
-class LatentApplyOperationCFG(io.ComfyNode):
+class LatentApplyOperationCFG:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentApplyOperationCFG",
-            category="latent/advanced/operations",
-            is_experimental=True,
-            inputs=[
-                io.Model.Input("model"),
-                io.LatentOperation.Input("operation"),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                             "operation": ("LATENT_OPERATION",),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, operation) -> io.NodeOutput:
+    CATEGORY = "latent/advanced/operations"
+    EXPERIMENTAL = True
+
+    def patch(self, model, operation):
         m = model.clone()
 
         def pre_cfg_function(args):
@@ -315,25 +261,21 @@ class LatentApplyOperationCFG(io.ComfyNode):
             return conds_out
 
         m.set_model_sampler_pre_cfg_function(pre_cfg_function)
-        return io.NodeOutput(m)
+        return (m, )
 
-class LatentOperationTonemapReinhard(io.ComfyNode):
+class LatentOperationTonemapReinhard:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentOperationTonemapReinhard",
-            category="latent/advanced/operations",
-            is_experimental=True,
-            inputs=[
-                io.Float.Input("multiplier", default=1.0, min=0.0, max=100.0, step=0.01),
-            ],
-            outputs=[
-                io.LatentOperation.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                              }}
 
-    @classmethod
-    def execute(cls, multiplier) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT_OPERATION",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced/operations"
+    EXPERIMENTAL = True
+
+    def op(self, multiplier):
         def tonemap_reinhard(latent, **kwargs):
             latent_vector_magnitude = (torch.linalg.vector_norm(latent, dim=(1)) + 0.0000000001)[:,None]
             normalized_latent = latent / latent_vector_magnitude
@@ -349,27 +291,39 @@ class LatentOperationTonemapReinhard(io.ComfyNode):
             new_magnitude *= top
 
             return normalized_latent * new_magnitude
-        return io.NodeOutput(tonemap_reinhard)
+        return (tonemap_reinhard,)
 
-class LatentOperationSharpen(io.ComfyNode):
+class LatentOperationSharpen:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LatentOperationSharpen",
-            category="latent/advanced/operations",
-            is_experimental=True,
-            inputs=[
-                io.Int.Input("sharpen_radius", default=9, min=1, max=31, step=1),
-                io.Float.Input("sigma", default=1.0, min=0.1, max=10.0, step=0.1),
-                io.Float.Input("alpha", default=0.1, min=0.0, max=5.0, step=0.01),
-            ],
-            outputs=[
-                io.LatentOperation.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+                "sharpen_radius": ("INT", {
+                    "default": 9,
+                    "min": 1,
+                    "max": 31,
+                    "step": 1
+                }),
+                "sigma": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 10.0,
+                    "step": 0.1
+                }),
+                "alpha": ("FLOAT", {
+                    "default": 0.1,
+                    "min": 0.0,
+                    "max": 5.0,
+                    "step": 0.01
+                }),
+                              }}
 
-    @classmethod
-    def execute(cls, sharpen_radius, sigma, alpha) -> io.NodeOutput:
+    RETURN_TYPES = ("LATENT_OPERATION",)
+    FUNCTION = "op"
+
+    CATEGORY = "latent/advanced/operations"
+    EXPERIMENTAL = True
+
+    def op(self, sharpen_radius, sigma, alpha):
         def sharpen(latent, **kwargs):
             luminance = (torch.linalg.vector_norm(latent, dim=(1)) + 1e-6)[:,None]
             normalized_latent = latent / luminance
@@ -386,27 +340,19 @@ class LatentOperationSharpen(io.ComfyNode):
             sharpened = torch.nn.functional.conv2d(padded_image, kernel.repeat(channels, 1, 1).unsqueeze(1), padding=kernel_size // 2, groups=channels)[:,:,sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius]
 
             return luminance * sharpened
-        return io.NodeOutput(sharpen)
+        return (sharpen,)
 
-
-class LatentExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            LatentAdd,
-            LatentSubtract,
-            LatentMultiply,
-            LatentInterpolate,
-            LatentConcat,
-            LatentCut,
-            LatentBatch,
-            LatentBatchSeedBehavior,
-            LatentApplyOperation,
-            LatentApplyOperationCFG,
-            LatentOperationTonemapReinhard,
-            LatentOperationSharpen,
-        ]
-
-
-async def comfy_entrypoint() -> LatentExtension:
-    return LatentExtension()
+NODE_CLASS_MAPPINGS = {
+    "LatentAdd": LatentAdd,
+    "LatentSubtract": LatentSubtract,
+    "LatentMultiply": LatentMultiply,
+    "LatentInterpolate": LatentInterpolate,
+    "LatentConcat": LatentConcat,
+    "LatentCut": LatentCut,
+    "LatentBatch": LatentBatch,
+    "LatentBatchSeedBehavior": LatentBatchSeedBehavior,
+    "LatentApplyOperation": LatentApplyOperation,
+    "LatentApplyOperationCFG": LatentApplyOperationCFG,
+    "LatentOperationTonemapReinhard": LatentOperationTonemapReinhard,
+    "LatentOperationSharpen": LatentOperationSharpen,
+}

comfy_extras/nodes_lora_extract.py

@@ -5,8 +5,6 @@ import folder_paths
 import os
 import logging
 from enum import Enum
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 CLAMP_QUANTILE = 0.99
 
@@ -73,40 +71,32 @@ def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, lora
             output_sd["{}{}.diff_b".format(prefix_lora, k[len(prefix_model):-5])] = sd[k].contiguous().half().cpu()
     return output_sd
 
-class LoraSave(io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LoraSave",
-            display_name="Extract and Save Lora",
-            category="_for_testing",
-            inputs=[
-                io.String.Input("filename_prefix", default="loras/ComfyUI_extracted_lora"),
-                io.Int.Input("rank", default=8, min=1, max=4096, step=1),
-                io.Combo.Input("lora_type", options=tuple(LORA_TYPES.keys())),
-                io.Boolean.Input("bias_diff", default=True),
-                io.Model.Input(
-                    "model_diff",
-                    tooltip="The ModelSubtract output to be converted to a lora.",
-                    optional=True,
-                ),
-                io.Clip.Input(
-                  "text_encoder_diff",
-                    tooltip="The CLIPSubtract output to be converted to a lora.",
-                    optional=True,
-                ),
-            ],
-            is_experimental=True,
-            is_output_node=True,
-        )
+class LoraSave:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
 
     @classmethod
-    def execute(cls, filename_prefix, rank, lora_type, bias_diff, model_diff=None, text_encoder_diff=None) -> io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {"required": {"filename_prefix": ("STRING", {"default": "loras/ComfyUI_extracted_lora"}),
+                              "rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1}),
+                              "lora_type": (tuple(LORA_TYPES.keys()),),
+                              "bias_diff": ("BOOLEAN", {"default": True}),
+                            },
+                "optional": {"model_diff": ("MODEL", {"tooltip": "The ModelSubtract output to be converted to a lora."}),
+                             "text_encoder_diff": ("CLIP", {"tooltip": "The CLIPSubtract output to be converted to a lora."})},
+    }
+    RETURN_TYPES = ()
+    FUNCTION = "save"
+    OUTPUT_NODE = True
+
+    CATEGORY = "_for_testing"
+
+    def save(self, filename_prefix, rank, lora_type, bias_diff, model_diff=None, text_encoder_diff=None):
         if model_diff is None and text_encoder_diff is None:
-            return io.NodeOutput()
+            return {}
 
         lora_type = LORA_TYPES.get(lora_type)
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
 
         output_sd = {}
         if model_diff is not None:
@@ -118,16 +108,12 @@ class LoraSave(io.ComfyNode):
         output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
 
         comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=None)
-        return io.NodeOutput()
+        return {}
 
+NODE_CLASS_MAPPINGS = {
+    "LoraSave": LoraSave
+}
 
-class LoraSaveExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            LoraSave,
-        ]
-
-
-async def comfy_entrypoint() -> LoraSaveExtension:
-    return LoraSaveExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LoraSave": "Extract and Save Lora"
+}

comfy_extras/nodes_lotus.py

@@ -1,22 +1,20 @@
-from typing_extensions import override
-
 import torch
 import comfy.model_management as mm
-from comfy_api.latest import ComfyExtension, io
 
-
-class LotusConditioning(io.ComfyNode):
+class LotusConditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LotusConditioning",
-            category="conditioning/lotus",
-            inputs=[],
-            outputs=[io.Conditioning.Output(display_name="conditioning")],
-        )
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+            },
+        }
 
-    @classmethod
-    def execute(cls) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    RETURN_NAMES = ("conditioning",)
+    FUNCTION = "conditioning"
+    CATEGORY = "conditioning/lotus"
+
+    def conditioning(self):
         device = mm.get_torch_device()
         #lotus uses a frozen encoder and null conditioning, i'm just inlining the results of that operation since it doesn't change
         #and getting parity with the reference implementation would otherwise require inference and 800mb of tensors
@@ -24,16 +22,8 @@ class LotusConditioning(io.ComfyNode):
 
         cond = [[prompt_embeds, {}]]
 
-        return io.NodeOutput(cond)
+        return (cond,)
 
-
-class LotusExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            LotusConditioning,
-        ]
-
-
-async def comfy_entrypoint() -> LotusExtension:
-    return LotusExtension()
+NODE_CLASS_MAPPINGS = {
+    "LotusConditioning" : LotusConditioning,
+}

comfy_extras/nodes_lt.py

@@ -1,3 +1,4 @@
+import io
 import nodes
 import node_helpers
 import torch
@@ -7,61 +8,46 @@ import comfy.utils
 import math
 import numpy as np
 import av
-from io import BytesIO
-from typing_extensions import override
 from comfy.ldm.lightricks.symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords
-from comfy_api.latest import ComfyExtension, io
 
-class EmptyLTXVLatentVideo(io.ComfyNode):
+class EmptyLTXVLatentVideo:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="EmptyLTXVLatentVideo",
-            category="latent/video/ltxv",
-            inputs=[
-                io.Int.Input("width", default=768, min=64, max=nodes.MAX_RESOLUTION, step=32),
-                io.Int.Input("height", default=512, min=64, max=nodes.MAX_RESOLUTION, step=32),
-                io.Int.Input("length", default=97, min=1, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "width": ("INT", {"default": 768, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 32}),
+                              "height": ("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 32}),
+                              "length": ("INT", {"default": 97, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "generate"
 
-    @classmethod
-    def execute(cls, width, height, length, batch_size=1) -> io.NodeOutput:
+    CATEGORY = "latent/video/ltxv"
+
+    def generate(self, width, height, length, batch_size=1):
         latent = torch.zeros([batch_size, 128, ((length - 1) // 8) + 1, height // 32, width // 32], device=comfy.model_management.intermediate_device())
-        return io.NodeOutput({"samples": latent})
+        return ({"samples": latent}, )
 
-    generate = execute  # TODO: remove
 
-class LTXVImgToVideo(io.ComfyNode):
+class LTXVImgToVideo:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVImgToVideo",
-            category="conditioning/video_models",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Vae.Input("vae"),
-                io.Image.Input("image"),
-                io.Int.Input("width", default=768, min=64, max=nodes.MAX_RESOLUTION, step=32),
-                io.Int.Input("height", default=512, min=64, max=nodes.MAX_RESOLUTION, step=32),
-                io.Int.Input("length", default=97, min=9, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-                io.Float.Input("strength", default=1.0, min=0.0, max=1.0),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "vae": ("VAE",),
+                             "image": ("IMAGE",),
+                             "width": ("INT", {"default": 768, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 32}),
+                             "height": ("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 32}),
+                             "length": ("INT", {"default": 97, "min": 9, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
+                             }}
 
-    @classmethod
-    def execute(cls, positive, negative, image, vae, width, height, length, batch_size, strength) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+
+    CATEGORY = "conditioning/video_models"
+    FUNCTION = "generate"
+
+    def generate(self, positive, negative, image, vae, width, height, length, batch_size, strength):
         pixels = comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
         encode_pixels = pixels[:, :, :, :3]
         t = vae.encode(encode_pixels)
@@ -76,9 +62,7 @@ class LTXVImgToVideo(io.ComfyNode):
         )
         conditioning_latent_frames_mask[:, :, :t.shape[2]] = 1.0 - strength
 
-        return io.NodeOutput(positive, negative, {"samples": latent, "noise_mask": conditioning_latent_frames_mask})
-
-    generate = execute  # TODO: remove
+        return (positive, negative, {"samples": latent, "noise_mask": conditioning_latent_frames_mask}, )
 
 
 def conditioning_get_any_value(conditioning, key, default=None):
@@ -109,46 +93,35 @@ def get_keyframe_idxs(cond):
     num_keyframes = torch.unique(keyframe_idxs[:, 0]).shape[0]
     return keyframe_idxs, num_keyframes
 
-class LTXVAddGuide(io.ComfyNode):
-    NUM_PREFIX_FRAMES = 2
-    PATCHIFIER = SymmetricPatchifier(1)
-
+class LTXVAddGuide:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVAddGuide",
-            category="conditioning/video_models",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Vae.Input("vae"),
-                io.Latent.Input("latent"),
-                io.Image.Input(
-                    "image",
-                    tooltip="Image or video to condition the latent video on. Must be 8*n + 1 frames. "
-                            "If the video is not 8*n + 1 frames, it will be cropped to the nearest 8*n + 1 frames.",
-                ),
-                io.Int.Input(
-                    "frame_idx",
-                    default=0,
-                    min=-9999,
-                    max=9999,
-                    tooltip="Frame index to start the conditioning at. "
-                            "For single-frame images or videos with 1-8 frames, any frame_idx value is acceptable. "
-                            "For videos with 9+ frames, frame_idx must be divisible by 8, otherwise it will be rounded "
-                            "down to the nearest multiple of 8. Negative values are counted from the end of the video.",
-                ),
-                io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "vae": ("VAE",),
+                             "latent": ("LATENT",),
+                             "image": ("IMAGE", {"tooltip": "Image or video to condition the latent video on. Must be 8*n + 1 frames."
+                                                 "If the video is not 8*n + 1 frames, it will be cropped to the nearest 8*n + 1 frames."}),
+                             "frame_idx": ("INT", {"default": 0, "min": -9999, "max": 9999,
+                                                   "tooltip": "Frame index to start the conditioning at. For single-frame images or "
+                                                   "videos with 1-8 frames, any frame_idx value is acceptable. For videos with 9+ "
+                                                   "frames, frame_idx must be divisible by 8, otherwise it will be rounded down to "
+                                                   "the nearest multiple of 8. Negative values are counted from the end of the video."}),
+                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                             }
+            }
 
-    @classmethod
-    def encode(cls, vae, latent_width, latent_height, images, scale_factors):
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+
+    CATEGORY = "conditioning/video_models"
+    FUNCTION = "generate"
+
+    def __init__(self):
+        self._num_prefix_frames = 2
+        self._patchifier = SymmetricPatchifier(1)
+
+    def encode(self, vae, latent_width, latent_height, images, scale_factors):
         time_scale_factor, width_scale_factor, height_scale_factor = scale_factors
         images = images[:(images.shape[0] - 1) // time_scale_factor * time_scale_factor + 1]
         pixels = comfy.utils.common_upscale(images.movedim(-1, 1), latent_width * width_scale_factor, latent_height * height_scale_factor, "bilinear", crop="disabled").movedim(1, -1)
@@ -156,8 +129,7 @@ class LTXVAddGuide(io.ComfyNode):
         t = vae.encode(encode_pixels)
         return encode_pixels, t
 
-    @classmethod
-    def get_latent_index(cls, cond, latent_length, guide_length, frame_idx, scale_factors):
+    def get_latent_index(self, cond, latent_length, guide_length, frame_idx, scale_factors):
         time_scale_factor, _, _ = scale_factors
         _, num_keyframes = get_keyframe_idxs(cond)
         latent_count = latent_length - num_keyframes
@@ -169,10 +141,9 @@ class LTXVAddGuide(io.ComfyNode):
 
         return frame_idx, latent_idx
 
-    @classmethod
-    def add_keyframe_index(cls, cond, frame_idx, guiding_latent, scale_factors):
+    def add_keyframe_index(self, cond, frame_idx, guiding_latent, scale_factors):
         keyframe_idxs, _ = get_keyframe_idxs(cond)
-        _, latent_coords = cls.PATCHIFIER.patchify(guiding_latent)
+        _, latent_coords = self._patchifier.patchify(guiding_latent)
         pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=frame_idx == 0)  # we need the causal fix only if we're placing the new latents at index 0
         pixel_coords[:, 0] += frame_idx
         if keyframe_idxs is None:
@@ -181,9 +152,8 @@ class LTXVAddGuide(io.ComfyNode):
             keyframe_idxs = torch.cat([keyframe_idxs, pixel_coords], dim=2)
         return node_helpers.conditioning_set_values(cond, {"keyframe_idxs": keyframe_idxs})
 
-    @classmethod
-    def append_keyframe(cls, positive, negative, frame_idx, latent_image, noise_mask, guiding_latent, strength, scale_factors):
-        _, latent_idx = cls.get_latent_index(
+    def append_keyframe(self, positive, negative, frame_idx, latent_image, noise_mask, guiding_latent, strength, scale_factors):
+        _, latent_idx = self.get_latent_index(
             cond=positive,
             latent_length=latent_image.shape[2],
             guide_length=guiding_latent.shape[2],
@@ -192,8 +162,8 @@ class LTXVAddGuide(io.ComfyNode):
         )
         noise_mask[:, :, latent_idx:latent_idx + guiding_latent.shape[2]] = 1.0
 
-        positive = cls.add_keyframe_index(positive, frame_idx, guiding_latent, scale_factors)
-        negative = cls.add_keyframe_index(negative, frame_idx, guiding_latent, scale_factors)
+        positive = self.add_keyframe_index(positive, frame_idx, guiding_latent, scale_factors)
+        negative = self.add_keyframe_index(negative, frame_idx, guiding_latent, scale_factors)
 
         mask = torch.full(
             (noise_mask.shape[0], 1, guiding_latent.shape[2], noise_mask.shape[3], noise_mask.shape[4]),
@@ -206,8 +176,7 @@ class LTXVAddGuide(io.ComfyNode):
         noise_mask = torch.cat([noise_mask, mask], dim=2)
         return positive, negative, latent_image, noise_mask
 
-    @classmethod
-    def replace_latent_frames(cls, latent_image, noise_mask, guiding_latent, latent_idx, strength):
+    def replace_latent_frames(self, latent_image, noise_mask, guiding_latent, latent_idx, strength):
         cond_length = guiding_latent.shape[2]
         assert latent_image.shape[2] >= latent_idx + cond_length, "Conditioning frames exceed the length of the latent sequence."
 
@@ -226,21 +195,20 @@ class LTXVAddGuide(io.ComfyNode):
 
         return latent_image, noise_mask
 
-    @classmethod
-    def execute(cls, positive, negative, vae, latent, image, frame_idx, strength) -> io.NodeOutput:
+    def generate(self, positive, negative, vae, latent, image, frame_idx, strength):
         scale_factors = vae.downscale_index_formula
         latent_image = latent["samples"]
         noise_mask = get_noise_mask(latent)
 
         _, _, latent_length, latent_height, latent_width = latent_image.shape
-        image, t = cls.encode(vae, latent_width, latent_height, image, scale_factors)
+        image, t = self.encode(vae, latent_width, latent_height, image, scale_factors)
 
-        frame_idx, latent_idx = cls.get_latent_index(positive, latent_length, len(image), frame_idx, scale_factors)
+        frame_idx, latent_idx = self.get_latent_index(positive, latent_length, len(image), frame_idx, scale_factors)
         assert latent_idx + t.shape[2] <= latent_length, "Conditioning frames exceed the length of the latent sequence."
 
-        num_prefix_frames = min(cls.NUM_PREFIX_FRAMES, t.shape[2])
+        num_prefix_frames = min(self._num_prefix_frames, t.shape[2])
 
-        positive, negative, latent_image, noise_mask = cls.append_keyframe(
+        positive, negative, latent_image, noise_mask = self.append_keyframe(
             positive,
             negative,
             frame_idx,
@@ -255,9 +223,9 @@ class LTXVAddGuide(io.ComfyNode):
 
         t = t[:, :, num_prefix_frames:]
         if t.shape[2] == 0:
-            return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask})
+            return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask},)
 
-        latent_image, noise_mask = cls.replace_latent_frames(
+        latent_image, noise_mask = self.replace_latent_frames(
             latent_image,
             noise_mask,
             t,
@@ -265,37 +233,34 @@ class LTXVAddGuide(io.ComfyNode):
             strength,
         )
 
-        return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask})
-
-    generate = execute  # TODO: remove
+        return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask},)
 
 
-class LTXVCropGuides(io.ComfyNode):
+class LTXVCropGuides:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVCropGuides",
-            category="conditioning/video_models",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Latent.Input("latent"),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "latent": ("LATENT",),
+                             }
+            }
 
-    @classmethod
-    def execute(cls, positive, negative, latent) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+
+    CATEGORY = "conditioning/video_models"
+    FUNCTION = "crop"
+
+    def __init__(self):
+        self._patchifier = SymmetricPatchifier(1)
+
+    def crop(self, positive, negative, latent):
         latent_image = latent["samples"].clone()
         noise_mask = get_noise_mask(latent)
 
         _, num_keyframes = get_keyframe_idxs(positive)
         if num_keyframes == 0:
-            return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask},)
+            return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask},)
 
         latent_image = latent_image[:, :, :-num_keyframes]
         noise_mask = noise_mask[:, :, :-num_keyframes]
@@ -303,54 +268,44 @@ class LTXVCropGuides(io.ComfyNode):
         positive = node_helpers.conditioning_set_values(positive, {"keyframe_idxs": None})
         negative = node_helpers.conditioning_set_values(negative, {"keyframe_idxs": None})
 
-        return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask})
-
-    crop = execute  # TODO: remove
+        return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask},)
 
 
-class LTXVConditioning(io.ComfyNode):
+class LTXVConditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVConditioning",
-            category="conditioning/video_models",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Float.Input("frame_rate", default=25.0, min=0.0, max=1000.0, step=0.01),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "frame_rate": ("FLOAT", {"default": 25.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
+    RETURN_NAMES = ("positive", "negative")
+    FUNCTION = "append"
 
-    @classmethod
-    def execute(cls, positive, negative, frame_rate) -> io.NodeOutput:
+    CATEGORY = "conditioning/video_models"
+
+    def append(self, positive, negative, frame_rate):
         positive = node_helpers.conditioning_set_values(positive, {"frame_rate": frame_rate})
         negative = node_helpers.conditioning_set_values(negative, {"frame_rate": frame_rate})
-        return io.NodeOutput(positive, negative)
+        return (positive, negative)
 
 
-class ModelSamplingLTXV(io.ComfyNode):
+class ModelSamplingLTXV:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ModelSamplingLTXV",
-            category="advanced/model",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("max_shift", default=2.05, min=0.0, max=100.0, step=0.01),
-                io.Float.Input("base_shift", default=0.95, min=0.0, max=100.0, step=0.01),
-                io.Latent.Input("latent", optional=True),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "max_shift": ("FLOAT", {"default": 2.05, "min": 0.0, "max": 100.0, "step":0.01}),
+                              "base_shift": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 100.0, "step":0.01}),
+                              },
+                "optional": {"latent": ("LATENT",), }
+                }
 
-    @classmethod
-    def execute(cls, model, max_shift, base_shift, latent=None) -> io.NodeOutput:
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "advanced/model"
+
+    def patch(self, model, max_shift, base_shift, latent=None):
         m = model.clone()
 
         if latent is None:
@@ -374,41 +329,37 @@ class ModelSamplingLTXV(io.ComfyNode):
         model_sampling.set_parameters(shift=shift)
         m.add_object_patch("model_sampling", model_sampling)
 
-        return io.NodeOutput(m)
+        return (m, )
 
 
-class LTXVScheduler(io.ComfyNode):
+class LTXVScheduler:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVScheduler",
-            category="sampling/custom_sampling/schedulers",
-            inputs=[
-                io.Int.Input("steps", default=20, min=1, max=10000),
-                io.Float.Input("max_shift", default=2.05, min=0.0, max=100.0, step=0.01),
-                io.Float.Input("base_shift", default=0.95, min=0.0, max=100.0, step=0.01),
-                io.Boolean.Input(
-                    id="stretch",
-                    default=True,
-                    tooltip="Stretch the sigmas to be in the range [terminal, 1].",
-                ),
-                io.Float.Input(
-                    id="terminal",
-                    default=0.1,
-                    min=0.0,
-                    max=0.99,
-                    step=0.01,
-                    tooltip="The terminal value of the sigmas after stretching.",
-                ),
-                io.Latent.Input("latent", optional=True),
-            ],
-            outputs=[
-                io.Sigmas.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                     "max_shift": ("FLOAT", {"default": 2.05, "min": 0.0, "max": 100.0, "step":0.01}),
+                     "base_shift": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 100.0, "step":0.01}),
+                     "stretch": ("BOOLEAN", {
+                        "default": True,
+                        "tooltip": "Stretch the sigmas to be in the range [terminal, 1]."
+                    }),
+                     "terminal": (
+                        "FLOAT",
+                        {
+                            "default": 0.1, "min": 0.0, "max": 0.99, "step": 0.01,
+                            "tooltip": "The terminal value of the sigmas after stretching."
+                        },
+                    ),
+                    },
+                "optional": {"latent": ("LATENT",), }
+               }
 
-    @classmethod
-    def execute(cls, steps, max_shift, base_shift, stretch, terminal, latent=None) -> io.NodeOutput:
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
+
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, steps, max_shift, base_shift, stretch, terminal, latent=None):
         if latent is None:
             tokens = 4096
         else:
@@ -438,7 +389,7 @@ class LTXVScheduler(io.ComfyNode):
             stretched = 1.0 - (one_minus_z / scale_factor)
             sigmas[non_zero_mask] = stretched
 
-        return io.NodeOutput(sigmas)
+        return (sigmas,)
 
 def encode_single_frame(output_file, image_array: np.ndarray, crf):
     container = av.open(output_file, "w", format="mp4")
@@ -472,55 +423,52 @@ def preprocess(image: torch.Tensor, crf=29):
         return image
 
     image_array = (image[:(image.shape[0] // 2) * 2, :(image.shape[1] // 2) * 2] * 255.0).byte().cpu().numpy()
-    with BytesIO() as output_file:
+    with io.BytesIO() as output_file:
         encode_single_frame(output_file, image_array, crf)
         video_bytes = output_file.getvalue()
-    with BytesIO(video_bytes) as video_file:
+    with io.BytesIO(video_bytes) as video_file:
         image_array = decode_single_frame(video_file)
     tensor = torch.tensor(image_array, dtype=image.dtype, device=image.device) / 255.0
     return tensor
 
 
-class LTXVPreprocess(io.ComfyNode):
+class LTXVPreprocess:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="LTXVPreprocess",
-            category="image",
-            inputs=[
-                io.Image.Input("image"),
-                io.Int.Input(
-                    id="img_compression", default=35, min=0, max=100, tooltip="Amount of compression to apply on image."
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "img_compression": (
+                    "INT",
+                    {
+                        "default": 35,
+                        "min": 0,
+                        "max": 100,
+                        "tooltip": "Amount of compression to apply on image.",
+                    },
                 ),
-            ],
-            outputs=[
-                io.Image.Output(display_name="output_image"),
-            ],
-        )
+            }
+        }
 
-    @classmethod
-    def execute(cls, image, img_compression) -> io.NodeOutput:
+    FUNCTION = "preprocess"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image",)
+    CATEGORY = "image"
+
+    def preprocess(self, image, img_compression):
         output_images = []
         for i in range(image.shape[0]):
             output_images.append(preprocess(image[i], img_compression))
-        return io.NodeOutput(torch.stack(output_images))
-
-    preprocess = execute  # TODO: remove
-
-class LtxvExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            EmptyLTXVLatentVideo,
-            LTXVImgToVideo,
-            ModelSamplingLTXV,
-            LTXVConditioning,
-            LTXVScheduler,
-            LTXVAddGuide,
-            LTXVPreprocess,
-            LTXVCropGuides,
-        ]
+        return (torch.stack(output_images),)
 
 
-async def comfy_entrypoint() -> LtxvExtension:
-    return LtxvExtension()
+NODE_CLASS_MAPPINGS = {
+    "EmptyLTXVLatentVideo": EmptyLTXVLatentVideo,
+    "LTXVImgToVideo": LTXVImgToVideo,
+    "ModelSamplingLTXV": ModelSamplingLTXV,
+    "LTXVConditioning": LTXVConditioning,
+    "LTXVScheduler": LTXVScheduler,
+    "LTXVAddGuide": LTXVAddGuide,
+    "LTXVPreprocess": LTXVPreprocess,
+    "LTXVCropGuides": LTXVCropGuides,
+}

comfy_extras/nodes_lumina2.py

@@ -1,27 +1,20 @@
-from typing_extensions import override
+from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
 import torch
 
-from comfy_api.latest import ComfyExtension, io
 
-
-class RenormCFG(io.ComfyNode):
+class RenormCFG:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="RenormCFG",
-            category="advanced/model",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("cfg_trunc", default=100, min=0.0, max=100.0, step=0.01),
-                io.Float.Input("renorm_cfg", default=1.0, min=0.0, max=100.0, step=0.01),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "cfg_trunc": ("FLOAT", {"default": 100, "min": 0.0, "max": 100.0, "step": 0.01}),
+                              "renorm_cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, cfg_trunc, renorm_cfg) -> io.NodeOutput:
+    CATEGORY = "advanced/model"
+
+    def patch(self, model, cfg_trunc, renorm_cfg):
         def renorm_cfg_func(args):
             cond_denoised = args["cond_denoised"]
             uncond_denoised = args["uncond_denoised"]
@@ -60,10 +53,10 @@ class RenormCFG(io.ComfyNode):
 
         m = model.clone()
         m.set_model_sampler_cfg_function(renorm_cfg_func)
-        return io.NodeOutput(m)
+        return (m, )
 
 
-class CLIPTextEncodeLumina2(io.ComfyNode):
+class CLIPTextEncodeLumina2(ComfyNodeABC):
     SYSTEM_PROMPT = {
         "superior": "You are an assistant designed to generate superior images with the superior "\
             "degree of image-text alignment based on textual prompts or user prompts.",
@@ -76,52 +69,36 @@ class CLIPTextEncodeLumina2(io.ComfyNode):
         "Alignment: You are an assistant designed to generate high-quality images with the highest "\
         "degree of image-text alignment based on textual prompts."
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeLumina2",
-            display_name="CLIP Text Encode for Lumina2",
-            category="conditioning",
-            description="Encodes a system prompt and a user prompt using a CLIP model into an embedding "
-                        "that can be used to guide the diffusion model towards generating specific images.",
-            inputs=[
-                io.Combo.Input(
-                    "system_prompt",
-                    options=list(cls.SYSTEM_PROMPT.keys()),
-                    tooltip=cls.SYSTEM_PROMPT_TIP,
-                ),
-                io.String.Input(
-                    "user_prompt",
-                    multiline=True,
-                    dynamic_prompts=True,
-                    tooltip="The text to be encoded.",
-                ),
-                io.Clip.Input("clip", tooltip="The CLIP model used for encoding the text."),
-            ],
-            outputs=[
-                io.Conditioning.Output(
-                    tooltip="A conditioning containing the embedded text used to guide the diffusion model.",
-                ),
-            ],
-        )
+    def INPUT_TYPES(s) -> InputTypeDict:
+        return {
+            "required": {
+                "system_prompt": (list(CLIPTextEncodeLumina2.SYSTEM_PROMPT.keys()), {"tooltip": CLIPTextEncodeLumina2.SYSTEM_PROMPT_TIP}),
+                "user_prompt": (IO.STRING, {"multiline": True, "dynamicPrompts": True, "tooltip": "The text to be encoded."}),
+                "clip": (IO.CLIP, {"tooltip": "The CLIP model used for encoding the text."})
+            }
+        }
+    RETURN_TYPES = (IO.CONDITIONING,)
+    OUTPUT_TOOLTIPS = ("A conditioning containing the embedded text used to guide the diffusion model.",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, clip, user_prompt, system_prompt) -> io.NodeOutput:
+    CATEGORY = "conditioning"
+    DESCRIPTION = "Encodes a system prompt and a user prompt using a CLIP model into an embedding that can be used to guide the diffusion model towards generating specific images."
+
+    def encode(self, clip, user_prompt, system_prompt):
         if clip is None:
             raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")
-        system_prompt = cls.SYSTEM_PROMPT[system_prompt]
+        system_prompt = CLIPTextEncodeLumina2.SYSTEM_PROMPT[system_prompt]
         prompt = f'{system_prompt} <Prompt Start> {user_prompt}'
         tokens = clip.tokenize(prompt)
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens))
+        return (clip.encode_from_tokens_scheduled(tokens), )
 
 
-class Lumina2Extension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            CLIPTextEncodeLumina2,
-            RenormCFG,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "CLIPTextEncodeLumina2": CLIPTextEncodeLumina2,
+    "RenormCFG": RenormCFG
+}
 
 
-async def comfy_entrypoint() -> Lumina2Extension:
-    return Lumina2Extension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CLIPTextEncodeLumina2": "CLIP Text Encode for Lumina2",
+}

comfy_extras/nodes_mahiro.py

@@ -1,29 +1,17 @@
-from typing_extensions import override
 import torch
 import torch.nn.functional as F
 
-from comfy_api.latest import ComfyExtension, io
-
-
-class Mahiro(io.ComfyNode):
+class Mahiro:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="Mahiro",
-            display_name="Mahiro is so cute that she deserves a better guidance function!! (。・ω・。)",
-            category="_for_testing",
-            description="Modify the guidance to scale more on the 'direction' of the positive prompt rather than the difference between the negative prompt.",
-            inputs=[
-                io.Model.Input("model"),
-            ],
-            outputs=[
-                io.Model.Output(display_name="patched_model"),
-            ],
-            is_experimental=True,
-        )
-
-    @classmethod
-    def execute(cls, model) -> io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL",),
+                            }}
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("patched_model",)
+    FUNCTION = "patch"
+    CATEGORY = "_for_testing"
+    DESCRIPTION = "Modify the guidance to scale more on the 'direction' of the positive prompt rather than the difference between the negative prompt."
+    def patch(self, model):
         m = model.clone()
         def mahiro_normd(args):
             scale: float = args['cond_scale']
@@ -42,16 +30,12 @@ class Mahiro(io.ComfyNode):
             wm = (simsc*cfg + (4-simsc)*leap) / 4
             return wm
         m.set_model_sampler_post_cfg_function(mahiro_normd)
-        return io.NodeOutput(m)
+        return (m, )
 
+NODE_CLASS_MAPPINGS = {
+    "Mahiro": Mahiro
+}
 
-class MahiroExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            Mahiro,
-        ]
-
-
-async def comfy_entrypoint() -> MahiroExtension:
-    return MahiroExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Mahiro": "Mahiro is so cute that she deserves a better guidance function!! (。・ω・。)",
+}

comfy_extras/nodes_mask.py

@@ -12,38 +12,35 @@ from nodes import MAX_RESOLUTION
 def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):
     source = source.to(destination.device)
     if resize_source:
-        source = torch.nn.functional.interpolate(source, size=(destination.shape[-2], destination.shape[-1]), mode="bilinear")
+        source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear")
 
     source = comfy.utils.repeat_to_batch_size(source, destination.shape[0])
 
-    x = max(-source.shape[-1] * multiplier, min(x, destination.shape[-1] * multiplier))
-    y = max(-source.shape[-2] * multiplier, min(y, destination.shape[-2] * multiplier))
+    x = max(-source.shape[3] * multiplier, min(x, destination.shape[3] * multiplier))
+    y = max(-source.shape[2] * multiplier, min(y, destination.shape[2] * multiplier))
 
     left, top = (x // multiplier, y // multiplier)
-    right, bottom = (left + source.shape[-1], top + source.shape[-2],)
+    right, bottom = (left + source.shape[3], top + source.shape[2],)
 
     if mask is None:
         mask = torch.ones_like(source)
     else:
         mask = mask.to(destination.device, copy=True)
-        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[-2], source.shape[-1]), mode="bilinear")
+        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
         mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])
 
     # calculate the bounds of the source that will be overlapping the destination
     # this prevents the source trying to overwrite latent pixels that are out of bounds
     # of the destination
-    visible_width, visible_height = (destination.shape[-1] - left + min(0, x), destination.shape[-2] - top + min(0, y),)
+    visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)
 
     mask = mask[:, :, :visible_height, :visible_width]
-    if mask.ndim < source.ndim:
-        mask = mask.unsqueeze(1)
-
     inverse_mask = torch.ones_like(mask) - mask
 
-    source_portion = mask * source[..., :visible_height, :visible_width]
-    destination_portion = inverse_mask  * destination[..., top:bottom, left:right]
+    source_portion = mask * source[:, :, :visible_height, :visible_width]
+    destination_portion = inverse_mask  * destination[:, :, top:bottom, left:right]
 
-    destination[..., top:bottom, left:right] = source_portion + destination_portion
+    destination[:, :, top:bottom, left:right] = source_portion + destination_portion
     return destination
 
 class LatentCompositeMasked:

comfy_extras/nodes_mochi.py

@@ -1,40 +1,23 @@
-from typing_extensions import override
+import nodes
 import torch
 import comfy.model_management
-import nodes
-from comfy_api.latest import ComfyExtension, io
 
-
-class EmptyMochiLatentVideo(io.ComfyNode):
+class EmptyMochiLatentVideo:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="EmptyMochiLatentVideo",
-            category="latent/video",
-            inputs=[
-                io.Int.Input("width", default=848, min=16, max=nodes.MAX_RESOLUTION, step=16),
-                io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
-                io.Int.Input("length", default=25, min=7, max=nodes.MAX_RESOLUTION, step=6),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "width": ("INT", {"default": 848, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                              "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                              "length": ("INT", {"default": 25, "min": 7, "max": nodes.MAX_RESOLUTION, "step": 6}),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "generate"
 
-    @classmethod
-    def execute(cls, width, height, length, batch_size=1) -> io.NodeOutput:
+    CATEGORY = "latent/video"
+
+    def generate(self, width, height, length, batch_size=1):
         latent = torch.zeros([batch_size, 12, ((length - 1) // 6) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
-        return io.NodeOutput({"samples": latent})
+        return ({"samples":latent}, )
 
-
-class MochiExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            EmptyMochiLatentVideo,
-        ]
-
-
-async def comfy_entrypoint() -> MochiExtension:
-    return MochiExtension()
+NODE_CLASS_MAPPINGS = {
+    "EmptyMochiLatentVideo": EmptyMochiLatentVideo,
+}

comfy_extras/nodes_model_downscale.py

@@ -1,33 +1,24 @@
-from typing_extensions import override
 import comfy.utils
-from comfy_api.latest import ComfyExtension, io
 
-
-class PatchModelAddDownscale(io.ComfyNode):
-    UPSCALE_METHODS = ["bicubic", "nearest-exact", "bilinear", "area", "bislerp"]
+class PatchModelAddDownscale:
+    upscale_methods = ["bicubic", "nearest-exact", "bilinear", "area", "bislerp"]
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PatchModelAddDownscale",
-            display_name="PatchModelAddDownscale (Kohya Deep Shrink)",
-            category="model_patches/unet",
-            inputs=[
-                io.Model.Input("model"),
-                io.Int.Input("block_number", default=3, min=1, max=32, step=1),
-                io.Float.Input("downscale_factor", default=2.0, min=0.1, max=9.0, step=0.001),
-                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("end_percent", default=0.35, min=0.0, max=1.0, step=0.001),
-                io.Boolean.Input("downscale_after_skip", default=True),
-                io.Combo.Input("downscale_method", options=cls.UPSCALE_METHODS),
-                io.Combo.Input("upscale_method", options=cls.UPSCALE_METHODS),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "block_number": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1}),
+                              "downscale_factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 9.0, "step": 0.001}),
+                              "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                              "end_percent": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}),
+                              "downscale_after_skip": ("BOOLEAN", {"default": True}),
+                              "downscale_method": (s.upscale_methods,),
+                              "upscale_method": (s.upscale_methods,),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip, downscale_method, upscale_method) -> io.NodeOutput:
+    CATEGORY = "model_patches/unet"
+
+    def patch(self, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip, downscale_method, upscale_method):
         model_sampling = model.get_model_object("model_sampling")
         sigma_start = model_sampling.percent_to_sigma(start_percent)
         sigma_end = model_sampling.percent_to_sigma(end_percent)
@@ -50,21 +41,13 @@ class PatchModelAddDownscale(io.ComfyNode):
         else:
             m.set_model_input_block_patch(input_block_patch)
         m.set_model_output_block_patch(output_block_patch)
-        return io.NodeOutput(m)
+        return (m, )
 
+NODE_CLASS_MAPPINGS = {
+    "PatchModelAddDownscale": PatchModelAddDownscale,
+}
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     # Sampling
-    "PatchModelAddDownscale": "",
+    "PatchModelAddDownscale": "PatchModelAddDownscale (Kohya Deep Shrink)",
 }
-
-class ModelDownscaleExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            PatchModelAddDownscale,
-        ]
-
-
-async def comfy_entrypoint() -> ModelDownscaleExtension:
-    return ModelDownscaleExtension()

comfy_extras/nodes_morphology.py

@@ -1,34 +1,24 @@
 import torch
 import comfy.model_management
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 from kornia.morphology import dilation, erosion, opening, closing, gradient, top_hat, bottom_hat
 import kornia.color
 
 
-class Morphology(io.ComfyNode):
+class Morphology:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="Morphology",
-            display_name="ImageMorphology",
-            category="image/postprocessing",
-            inputs=[
-                io.Image.Input("image"),
-                io.Combo.Input(
-                    "operation",
-                    options=["erode", "dilate", "open", "close", "gradient", "bottom_hat", "top_hat"],
-                ),
-                io.Int.Input("kernel_size", default=3, min=3, max=999, step=1),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"image": ("IMAGE",),
+                                "operation": (["erode",  "dilate", "open", "close", "gradient", "bottom_hat", "top_hat"],),
+                                "kernel_size": ("INT", {"default": 3, "min": 3, "max": 999, "step": 1}),
+                                }}
 
-    @classmethod
-    def execute(cls, image, operation, kernel_size) -> io.NodeOutput:
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "process"
+
+    CATEGORY = "image/postprocessing"
+
+    def process(self, image, operation, kernel_size):
         device = comfy.model_management.get_torch_device()
         kernel = torch.ones(kernel_size, kernel_size, device=device)
         image_k = image.to(device).movedim(-1, 1)
@@ -49,63 +39,49 @@ class Morphology(io.ComfyNode):
         else:
             raise ValueError(f"Invalid operation {operation} for morphology. Must be one of 'erode', 'dilate', 'open', 'close', 'gradient', 'tophat', 'bottomhat'")
         img_out = output.to(comfy.model_management.intermediate_device()).movedim(1, -1)
-        return io.NodeOutput(img_out)
+        return (img_out,)
 
 
-class ImageRGBToYUV(io.ComfyNode):
+class ImageRGBToYUV:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageRGBToYUV",
-            category="image/batch",
-            inputs=[
-                io.Image.Input("image"),
-            ],
-            outputs=[
-                io.Image.Output(display_name="Y"),
-                io.Image.Output(display_name="U"),
-                io.Image.Output(display_name="V"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "image": ("IMAGE",),
+                              }}
 
-    @classmethod
-    def execute(cls, image) -> io.NodeOutput:
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
+    RETURN_NAMES = ("Y", "U", "V")
+    FUNCTION = "execute"
+
+    CATEGORY = "image/batch"
+
+    def execute(self, image):
         out = kornia.color.rgb_to_ycbcr(image.movedim(-1, 1)).movedim(1, -1)
-        return io.NodeOutput(out[..., 0:1].expand_as(image), out[..., 1:2].expand_as(image), out[..., 2:3].expand_as(image))
+        return (out[..., 0:1].expand_as(image), out[..., 1:2].expand_as(image), out[..., 2:3].expand_as(image))
 
-class ImageYUVToRGB(io.ComfyNode):
+class ImageYUVToRGB:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageYUVToRGB",
-            category="image/batch",
-            inputs=[
-                io.Image.Input("Y"),
-                io.Image.Input("U"),
-                io.Image.Input("V"),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"Y": ("IMAGE",),
+                             "U": ("IMAGE",),
+                             "V": ("IMAGE",),
+                              }}
 
-    @classmethod
-    def execute(cls, Y, U, V) -> io.NodeOutput:
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "image/batch"
+
+    def execute(self, Y, U, V):
         image = torch.cat([torch.mean(Y, dim=-1, keepdim=True), torch.mean(U, dim=-1, keepdim=True), torch.mean(V, dim=-1, keepdim=True)], dim=-1)
         out = kornia.color.ycbcr_to_rgb(image.movedim(-1, 1)).movedim(1, -1)
-        return io.NodeOutput(out)
+        return (out,)
 
+NODE_CLASS_MAPPINGS = {
+    "Morphology": Morphology,
+    "ImageRGBToYUV": ImageRGBToYUV,
+    "ImageYUVToRGB": ImageYUVToRGB,
+}
 
-class MorphologyExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            Morphology,
-            ImageRGBToYUV,
-            ImageYUVToRGB,
-        ]
-
-
-async def comfy_entrypoint() -> MorphologyExtension:
-    return MorphologyExtension()
-
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Morphology": "ImageMorphology",
+}

comfy_extras/nodes_optimalsteps.py

@@ -1,12 +1,9 @@
 # from https://github.com/bebebe666/OptimalSteps
 
+
 import numpy as np
 import torch
 
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
-
-
 def loglinear_interp(t_steps, num_steps):
     """
     Performs log-linear interpolation of a given array of decreasing numbers.
@@ -26,28 +23,25 @@ NOISE_LEVELS = {"FLUX": [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0
 "Chroma": [0.992, 0.99, 0.988, 0.985, 0.982, 0.978, 0.973, 0.968, 0.961, 0.953, 0.943, 0.931, 0.917, 0.9, 0.881, 0.858, 0.832, 0.802, 0.769, 0.731, 0.69, 0.646, 0.599, 0.55, 0.501, 0.451, 0.402, 0.355, 0.311, 0.27, 0.232, 0.199, 0.169, 0.143, 0.12, 0.101, 0.084, 0.07, 0.058, 0.048, 0.001],
 }
 
-class OptimalStepsScheduler(io.ComfyNode):
+class OptimalStepsScheduler:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="OptimalStepsScheduler",
-            category="sampling/custom_sampling/schedulers",
-            inputs=[
-                io.Combo.Input("model_type", options=["FLUX", "Wan", "Chroma"]),
-                io.Int.Input("steps", default=20, min=3, max=1000),
-                io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
-            ],
-            outputs=[
-                io.Sigmas.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model_type": (["FLUX", "Wan", "Chroma"], ),
+                     "steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
+                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
 
-    @classmethod
-    def execute(cls, model_type, steps, denoise) ->io.NodeOutput:
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, model_type, steps, denoise):
         total_steps = steps
         if denoise < 1.0:
             if denoise <= 0.0:
-                return io.NodeOutput(torch.FloatTensor([]))
+                return (torch.FloatTensor([]),)
             total_steps = round(steps * denoise)
 
         sigmas = NOISE_LEVELS[model_type][:]
@@ -56,16 +50,8 @@ class OptimalStepsScheduler(io.ComfyNode):
 
         sigmas = sigmas[-(total_steps + 1):]
         sigmas[-1] = 0
-        return io.NodeOutput(torch.FloatTensor(sigmas))
+        return (torch.FloatTensor(sigmas), )
 
-
-class OptimalStepsExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            OptimalStepsScheduler,
-        ]
-
-
-async def comfy_entrypoint() -> OptimalStepsExtension:
-    return OptimalStepsExtension()
+NODE_CLASS_MAPPINGS = {
+    "OptimalStepsScheduler": OptimalStepsScheduler,
+}

comfy_extras/nodes_pag.py

@@ -3,30 +3,25 @@
 
 #My modified one here is more basic but has less chances of breaking with ComfyUI updates.
 
-from typing_extensions import override
-
 import comfy.model_patcher
 import comfy.samplers
-from comfy_api.latest import ComfyExtension, io
 
-
-class PerturbedAttentionGuidance(io.ComfyNode):
+class PerturbedAttentionGuidance:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PerturbedAttentionGuidance",
-            category="model_patches/unet",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("scale", default=3.0, min=0.0, max=100.0, step=0.01, round=0.01),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.01, "round": 0.01}),
+            }
+        }
 
-    @classmethod
-    def execute(cls, model, scale) -> io.NodeOutput:
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "model_patches/unet"
+
+    def patch(self, model, scale):
         unet_block = "middle"
         unet_block_id = 0
         m = model.clone()
@@ -54,16 +49,8 @@ class PerturbedAttentionGuidance(io.ComfyNode):
 
         m.set_model_sampler_post_cfg_function(post_cfg_function)
 
-        return io.NodeOutput(m)
+        return (m,)
 
-
-class PAGExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            PerturbedAttentionGuidance,
-        ]
-
-
-async def comfy_entrypoint() -> PAGExtension:
-    return PAGExtension()
+NODE_CLASS_MAPPINGS = {
+    "PerturbedAttentionGuidance": PerturbedAttentionGuidance,
+}

comfy_extras/nodes_perpneg.py

@@ -5,9 +5,6 @@ import comfy.samplers
 import comfy.utils
 import node_helpers
 import math
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
-
 
 def perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, cond_scale):
     pos = noise_pred_pos - noise_pred_nocond
@@ -19,27 +16,20 @@ def perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, co
     return cfg_result
 
 #TODO: This node should be removed, it has been replaced with PerpNegGuider
-class PerpNeg(io.ComfyNode):
+class PerpNeg:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PerpNeg",
-            display_name="Perp-Neg (DEPRECATED by PerpNegGuider)",
-            category="_for_testing",
-            inputs=[
-                io.Model.Input("model"),
-                io.Conditioning.Input("empty_conditioning"),
-                io.Float.Input("neg_scale", default=1.0, min=0.0, max=100.0, step=0.01),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-            is_experimental=True,
-            is_deprecated=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL", ),
+                             "empty_conditioning": ("CONDITIONING", ),
+                             "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                            }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, empty_conditioning, neg_scale) -> io.NodeOutput:
+    CATEGORY = "_for_testing"
+    DEPRECATED = True
+
+    def patch(self, model, empty_conditioning, neg_scale):
         m = model.clone()
         nocond = comfy.sampler_helpers.convert_cond(empty_conditioning)
 
@@ -60,7 +50,7 @@ class PerpNeg(io.ComfyNode):
 
         m.set_model_sampler_cfg_function(cfg_function)
 
-        return io.NodeOutput(m)
+        return (m, )
 
 
 class Guider_PerpNeg(comfy.samplers.CFGGuider):
@@ -122,42 +112,35 @@ class Guider_PerpNeg(comfy.samplers.CFGGuider):
 
         return cfg_result
 
-class PerpNegGuider(io.ComfyNode):
+class PerpNegGuider:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PerpNegGuider",
-            category="_for_testing",
-            inputs=[
-                io.Model.Input("model"),
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Conditioning.Input("empty_conditioning"),
-                io.Float.Input("cfg", default=8.0, min=0.0, max=100.0, step=0.1, round=0.01),
-                io.Float.Input("neg_scale", default=1.0, min=0.0, max=100.0, step=0.01),
-            ],
-            outputs=[
-                io.Guider.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                    "positive": ("CONDITIONING", ),
+                    "negative": ("CONDITIONING", ),
+                    "empty_conditioning": ("CONDITIONING", ),
+                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                    "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                     }
+                }
 
-    @classmethod
-    def execute(cls, model, positive, negative, empty_conditioning, cfg, neg_scale) -> io.NodeOutput:
+    RETURN_TYPES = ("GUIDER",)
+
+    FUNCTION = "get_guider"
+    CATEGORY = "_for_testing"
+
+    def get_guider(self, model, positive, negative, empty_conditioning, cfg, neg_scale):
         guider = Guider_PerpNeg(model)
         guider.set_conds(positive, negative, empty_conditioning)
         guider.set_cfg(cfg, neg_scale)
-        return io.NodeOutput(guider)
+        return (guider,)
 
+NODE_CLASS_MAPPINGS = {
+    "PerpNeg": PerpNeg,
+    "PerpNegGuider": PerpNegGuider,
+}
 
-class PerpNegExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            PerpNeg,
-            PerpNegGuider,
-        ]
-
-
-async def comfy_entrypoint() -> PerpNegExtension:
-    return PerpNegExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PerpNeg": "Perp-Neg (DEPRECATED by PerpNegGuider)",
+}

comfy_extras/nodes_photomaker.py

@@ -4,8 +4,6 @@ import folder_paths
 import comfy.clip_model
 import comfy.clip_vision
 import comfy.ops
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 # code for model from: https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model.py under Apache License Version 2.0
 VISION_CONFIG_DICT = {
@@ -118,52 +116,41 @@ class PhotoMakerIDEncoder(comfy.clip_model.CLIPVisionModelProjection):
         return updated_prompt_embeds
 
 
-class PhotoMakerLoader(io.ComfyNode):
+class PhotoMakerLoader:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PhotoMakerLoader",
-            category="_for_testing/photomaker",
-            inputs=[
-                io.Combo.Input("photomaker_model_name", options=folder_paths.get_filename_list("photomaker")),
-            ],
-            outputs=[
-                io.Photomaker.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "photomaker_model_name": (folder_paths.get_filename_list("photomaker"), )}}
 
-    @classmethod
-    def execute(cls, photomaker_model_name):
+    RETURN_TYPES = ("PHOTOMAKER",)
+    FUNCTION = "load_photomaker_model"
+
+    CATEGORY = "_for_testing/photomaker"
+
+    def load_photomaker_model(self, photomaker_model_name):
         photomaker_model_path = folder_paths.get_full_path_or_raise("photomaker", photomaker_model_name)
         photomaker_model = PhotoMakerIDEncoder()
         data = comfy.utils.load_torch_file(photomaker_model_path, safe_load=True)
         if "id_encoder" in data:
             data = data["id_encoder"]
         photomaker_model.load_state_dict(data)
-        return io.NodeOutput(photomaker_model)
+        return (photomaker_model,)
 
 
-class PhotoMakerEncode(io.ComfyNode):
+class PhotoMakerEncode:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="PhotoMakerEncode",
-            category="_for_testing/photomaker",
-            inputs=[
-                io.Photomaker.Input("photomaker"),
-                io.Image.Input("image"),
-                io.Clip.Input("clip"),
-                io.String.Input("text", multiline=True, dynamic_prompts=True, default="photograph of photomaker"),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "photomaker": ("PHOTOMAKER",),
+                              "image": ("IMAGE",),
+                              "clip": ("CLIP", ),
+                              "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "photograph of photomaker"}),
+                             }}
 
-    @classmethod
-    def execute(cls, photomaker, image, clip, text):
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "apply_photomaker"
+
+    CATEGORY = "_for_testing/photomaker"
+
+    def apply_photomaker(self, photomaker, image, clip, text):
         special_token = "photomaker"
         pixel_values = comfy.clip_vision.clip_preprocess(image.to(photomaker.load_device)).float()
         try:
@@ -191,16 +178,11 @@ class PhotoMakerEncode(io.ComfyNode):
         else:
             out = cond
 
-        return io.NodeOutput([[out, {"pooled_output": pooled}]])
+        return ([[out, {"pooled_output": pooled}]], )
 
 
-class PhotomakerExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            PhotoMakerLoader,
-            PhotoMakerEncode,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "PhotoMakerLoader": PhotoMakerLoader,
+    "PhotoMakerEncode": PhotoMakerEncode,
+}
 
-async def comfy_entrypoint() -> PhotomakerExtension:
-    return PhotomakerExtension()

comfy_extras/nodes_pixart.py

@@ -1,38 +1,24 @@
-from typing_extensions import override
-import nodes
-from comfy_api.latest import ComfyExtension, io
+from nodes import MAX_RESOLUTION
 
-class CLIPTextEncodePixArtAlpha(io.ComfyNode):
+class CLIPTextEncodePixArtAlpha:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodePixArtAlpha",
-            category="advanced/conditioning",
-            description="Encodes text and sets the resolution conditioning for PixArt Alpha. Does not apply to PixArt Sigma.",
-            inputs=[
-                io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
-                # "aspect_ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                io.String.Input("text", multiline=True, dynamic_prompts=True),
-                io.Clip.Input("clip"),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            # "aspect_ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+            "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ),
+            }}
 
-    @classmethod
-    def execute(cls, clip, width, height, text):
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+    CATEGORY = "advanced/conditioning"
+    DESCRIPTION = "Encodes text and sets the resolution conditioning for PixArt Alpha. Does not apply to PixArt Sigma."
+
+    def encode(self, clip, width, height, text):
         tokens = clip.tokenize(text)
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height}))
+        return (clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height}),)
 
-
-class PixArtExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            CLIPTextEncodePixArtAlpha,
-        ]
-
-async def comfy_entrypoint() -> PixArtExtension:
-    return PixArtExtension()
+NODE_CLASS_MAPPINGS = {
+    "CLIPTextEncodePixArtAlpha": CLIPTextEncodePixArtAlpha,
+}

comfy_extras/nodes_post_processing.py

@@ -1,4 +1,3 @@
-from typing_extensions import override
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -8,27 +7,33 @@ import math
 import comfy.utils
 import comfy.model_management
 import node_helpers
-from comfy_api.latest import ComfyExtension, io
 
-class Blend(io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageBlend",
-            category="image/postprocessing",
-            inputs=[
-                io.Image.Input("image1"),
-                io.Image.Input("image2"),
-                io.Float.Input("blend_factor", default=0.5, min=0.0, max=1.0, step=0.01),
-                io.Combo.Input("blend_mode", options=["normal", "multiply", "screen", "overlay", "soft_light", "difference"]),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+class Blend:
+    def __init__(self):
+        pass
 
     @classmethod
-    def execute(cls, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str) -> io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image1": ("IMAGE",),
+                "image2": ("IMAGE",),
+                "blend_factor": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light", "difference"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "blend_images"
+
+    CATEGORY = "image/postprocessing"
+
+    def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str):
         image1, image2 = node_helpers.image_alpha_fix(image1, image2)
         image2 = image2.to(image1.device)
         if image1.shape != image2.shape:
@@ -36,13 +41,12 @@ class Blend(io.ComfyNode):
             image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')
             image2 = image2.permute(0, 2, 3, 1)
 
-        blended_image = cls.blend_mode(image1, image2, blend_mode)
+        blended_image = self.blend_mode(image1, image2, blend_mode)
         blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
         blended_image = torch.clamp(blended_image, 0, 1)
-        return io.NodeOutput(blended_image)
+        return (blended_image,)
 
-    @classmethod
-    def blend_mode(cls, img1, img2, mode):
+    def blend_mode(self, img1, img2, mode):
         if mode == "normal":
             return img2
         elif mode == "multiply":
@@ -52,13 +56,13 @@ class Blend(io.ComfyNode):
         elif mode == "overlay":
             return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
         elif mode == "soft_light":
-            return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (cls.g(img1) - img1))
+            return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1))
         elif mode == "difference":
             return img1 - img2
-        raise ValueError(f"Unsupported blend mode: {mode}")
+        else:
+            raise ValueError(f"Unsupported blend mode: {mode}")
 
-    @classmethod
-    def g(cls, x):
+    def g(self, x):
         return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
 
 def gaussian_kernel(kernel_size: int, sigma: float, device=None):
@@ -67,26 +71,38 @@ def gaussian_kernel(kernel_size: int, sigma: float, device=None):
     g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
     return g / g.sum()
 
-class Blur(io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageBlur",
-            category="image/postprocessing",
-            inputs=[
-                io.Image.Input("image"),
-                io.Int.Input("blur_radius", default=1, min=1, max=31, step=1),
-                io.Float.Input("sigma", default=1.0, min=0.1, max=10.0, step=0.1),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+class Blur:
+    def __init__(self):
+        pass
 
     @classmethod
-    def execute(cls, image: torch.Tensor, blur_radius: int, sigma: float) -> io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "blur_radius": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 31,
+                    "step": 1
+                }),
+                "sigma": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 10.0,
+                    "step": 0.1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "blur"
+
+    CATEGORY = "image/postprocessing"
+
+    def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
         if blur_radius == 0:
-            return io.NodeOutput(image)
+            return (image,)
 
         image = image.to(comfy.model_management.get_torch_device())
         batch_size, height, width, channels = image.shape
@@ -99,24 +115,31 @@ class Blur(io.ComfyNode):
         blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius]
         blurred = blurred.permute(0, 2, 3, 1)
 
-        return io.NodeOutput(blurred.to(comfy.model_management.intermediate_device()))
+        return (blurred.to(comfy.model_management.intermediate_device()),)
 
+class Quantize:
+    def __init__(self):
+        pass
 
-class Quantize(io.ComfyNode):
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageQuantize",
-            category="image/postprocessing",
-            inputs=[
-                io.Image.Input("image"),
-                io.Int.Input("colors", default=256, min=1, max=256, step=1),
-                io.Combo.Input("dither", options=["none", "floyd-steinberg", "bayer-2", "bayer-4", "bayer-8", "bayer-16"]),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "colors": ("INT", {
+                    "default": 256,
+                    "min": 1,
+                    "max": 256,
+                    "step": 1
+                }),
+                "dither": (["none", "floyd-steinberg", "bayer-2", "bayer-4", "bayer-8", "bayer-16"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "quantize"
+
+    CATEGORY = "image/postprocessing"
 
     @staticmethod
     def bayer(im, pal_im, order):
@@ -144,8 +167,7 @@ class Quantize(io.ComfyNode):
         im = im.quantize(palette=pal_im, dither=Image.Dither.NONE)
         return im
 
-    @classmethod
-    def execute(cls, image: torch.Tensor, colors: int, dither: str) -> io.NodeOutput:
+    def quantize(self, image: torch.Tensor, colors: int, dither: str):
         batch_size, height, width, _ = image.shape
         result = torch.zeros_like(image)
 
@@ -165,29 +187,46 @@ class Quantize(io.ComfyNode):
             quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255
             result[b] = quantized_array
 
-        return io.NodeOutput(result)
+        return (result,)
 
-class Sharpen(io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageSharpen",
-            category="image/postprocessing",
-            inputs=[
-                io.Image.Input("image"),
-                io.Int.Input("sharpen_radius", default=1, min=1, max=31, step=1),
-                io.Float.Input("sigma", default=1.0, min=0.1, max=10.0, step=0.01),
-                io.Float.Input("alpha", default=1.0, min=0.0, max=5.0, step=0.01),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+class Sharpen:
+    def __init__(self):
+        pass
 
     @classmethod
-    def execute(cls, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha: float) -> io.NodeOutput:
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "sharpen_radius": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 31,
+                    "step": 1
+                }),
+                "sigma": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 10.0,
+                    "step": 0.01
+                }),
+                "alpha": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 5.0,
+                    "step": 0.01
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "sharpen"
+
+    CATEGORY = "image/postprocessing"
+
+    def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha: float):
         if sharpen_radius == 0:
-            return io.NodeOutput(image)
+            return (image,)
 
         batch_size, height, width, channels = image.shape
         image = image.to(comfy.model_management.get_torch_device())
@@ -206,29 +245,23 @@ class Sharpen(io.ComfyNode):
 
         result = torch.clamp(sharpened, 0, 1)
 
-        return io.NodeOutput(result.to(comfy.model_management.intermediate_device()))
+        return (result.to(comfy.model_management.intermediate_device()),)
 
-class ImageScaleToTotalPixels(io.ComfyNode):
+class ImageScaleToTotalPixels:
     upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
     crop_methods = ["disabled", "center"]
 
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="ImageScaleToTotalPixels",
-            category="image/upscaling",
-            inputs=[
-                io.Image.Input("image"),
-                io.Combo.Input("upscale_method", options=cls.upscale_methods),
-                io.Float.Input("megapixels", default=1.0, min=0.01, max=16.0, step=0.01),
-            ],
-            outputs=[
-                io.Image.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
+                              "megapixels": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01}),
+                            }}
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "upscale"
 
-    @classmethod
-    def execute(cls, image, upscale_method, megapixels) -> io.NodeOutput:
+    CATEGORY = "image/upscaling"
+
+    def upscale(self, image, upscale_method, megapixels):
         samples = image.movedim(-1,1)
         total = int(megapixels * 1024 * 1024)
 
@@ -238,18 +271,12 @@ class ImageScaleToTotalPixels(io.ComfyNode):
 
         s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
         s = s.movedim(1,-1)
-        return io.NodeOutput(s)
+        return (s,)
 
-class PostProcessingExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            Blend,
-            Blur,
-            Quantize,
-            Sharpen,
-            ImageScaleToTotalPixels,
-        ]
-
-async def comfy_entrypoint() -> PostProcessingExtension:
-    return PostProcessingExtension()
+NODE_CLASS_MAPPINGS = {
+    "ImageBlend": Blend,
+    "ImageBlur": Blur,
+    "ImageQuantize": Quantize,
+    "ImageSharpen": Sharpen,
+    "ImageScaleToTotalPixels": ImageScaleToTotalPixels,
+}

comfy_extras/nodes_qwen.py

@@ -1,29 +1,24 @@
 import node_helpers
 import comfy.utils
 import math
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
-class TextEncodeQwenImageEdit(io.ComfyNode):
+class TextEncodeQwenImageEdit:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="TextEncodeQwenImageEdit",
-            category="advanced/conditioning",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
-                io.Vae.Input("vae", optional=True),
-                io.Image.Input("image", optional=True),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            },
+            "optional": {"vae": ("VAE", ),
+                         "image": ("IMAGE", ),}}
 
-    @classmethod
-    def execute(cls, clip, prompt, vae=None, image=None) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, prompt, vae=None, image=None):
         ref_latent = None
         if image is None:
             images = []
@@ -45,30 +40,28 @@ class TextEncodeQwenImageEdit(io.ComfyNode):
         conditioning = clip.encode_from_tokens_scheduled(tokens)
         if ref_latent is not None:
             conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": [ref_latent]}, append=True)
-        return io.NodeOutput(conditioning)
+        return (conditioning, )
 
 
-class TextEncodeQwenImageEditPlus(io.ComfyNode):
+class TextEncodeQwenImageEditPlus:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="TextEncodeQwenImageEditPlus",
-            category="advanced/conditioning",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
-                io.Vae.Input("vae", optional=True),
-                io.Image.Input("image1", optional=True),
-                io.Image.Input("image2", optional=True),
-                io.Image.Input("image3", optional=True),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            },
+            "optional": {"vae": ("VAE", ),
+                         "image1": ("IMAGE", ),
+                         "image2": ("IMAGE", ),
+                         "image3": ("IMAGE", ),
+                         }}
 
-    @classmethod
-    def execute(cls, clip, prompt, vae=None, image1=None, image2=None, image3=None) -> io.NodeOutput:
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, prompt, vae=None, image1=None, image2=None, image3=None):
         ref_latents = []
         images = [image1, image2, image3]
         images_vl = []
@@ -101,17 +94,10 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode):
         conditioning = clip.encode_from_tokens_scheduled(tokens)
         if len(ref_latents) > 0:
             conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True)
-        return io.NodeOutput(conditioning)
+        return (conditioning, )
 
 
-class QwenExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            TextEncodeQwenImageEdit,
-            TextEncodeQwenImageEditPlus,
-        ]
-
-
-async def comfy_entrypoint() -> QwenExtension:
-    return QwenExtension()
+NODE_CLASS_MAPPINGS = {
+    "TextEncodeQwenImageEdit": TextEncodeQwenImageEdit,
+    "TextEncodeQwenImageEditPlus": TextEncodeQwenImageEditPlus,
+}

comfy_extras/nodes_rebatch.py

@@ -1,25 +1,18 @@
-from typing_extensions import override
 import torch
 
-from comfy_api.latest import ComfyExtension, io
-
-
-class LatentRebatch(io.ComfyNode):
+class LatentRebatch:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="RebatchLatents",
-            display_name="Rebatch Latents",
-            category="latent/batch",
-            is_input_list=True,
-            inputs=[
-                io.Latent.Input("latents"),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-            ],
-            outputs=[
-                io.Latent.Output(is_output_list=True),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "latents": ("LATENT",),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                              }}
+    RETURN_TYPES = ("LATENT",)
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True, )
+
+    FUNCTION = "rebatch"
+
+    CATEGORY = "latent/batch"
 
     @staticmethod
     def get_batch(latents, list_ind, offset):
@@ -60,8 +53,7 @@ class LatentRebatch(io.ComfyNode):
         result = [torch.cat((b1, b2)) if torch.is_tensor(b1) else b1 + b2 for b1, b2 in zip(batch1, batch2)]
         return result
 
-    @classmethod
-    def execute(cls, latents, batch_size):
+    def rebatch(self, latents, batch_size):
         batch_size = batch_size[0]
 
         output_list = []
@@ -71,24 +63,24 @@ class LatentRebatch(io.ComfyNode):
         for i in range(len(latents)):
             # fetch new entry of list
             #samples, masks, indices = self.get_batch(latents, i)
-            next_batch = cls.get_batch(latents, i, processed)
+            next_batch = self.get_batch(latents, i, processed)
             processed += len(next_batch[2])
             # set to current if current is None
             if current_batch[0] is None:
                 current_batch = next_batch
             # add previous to list if dimensions do not match
             elif next_batch[0].shape[-1] != current_batch[0].shape[-1] or next_batch[0].shape[-2] != current_batch[0].shape[-2]:
-                sliced, _ = cls.slice_batch(current_batch, 1, batch_size)
+                sliced, _ = self.slice_batch(current_batch, 1, batch_size)
                 output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
                 current_batch = next_batch
             # cat if everything checks out
             else:
-                current_batch = cls.cat_batch(current_batch, next_batch)
+                current_batch = self.cat_batch(current_batch, next_batch)
 
             # add to list if dimensions gone above target batch size
             if current_batch[0].shape[0] > batch_size:
                 num = current_batch[0].shape[0] // batch_size
-                sliced, remainder = cls.slice_batch(current_batch, num, batch_size)
+                sliced, remainder = self.slice_batch(current_batch, num, batch_size)
 
                 for i in range(num):
                     output_list.append({'samples': sliced[0][i], 'noise_mask': sliced[1][i], 'batch_index': sliced[2][i]})
@@ -97,7 +89,7 @@ class LatentRebatch(io.ComfyNode):
 
         #add remainder
         if current_batch[0] is not None:
-            sliced, _ = cls.slice_batch(current_batch, 1, batch_size)
+            sliced, _ = self.slice_batch(current_batch, 1, batch_size)
             output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
 
         #get rid of empty masks
@@ -105,27 +97,23 @@ class LatentRebatch(io.ComfyNode):
             if s['noise_mask'].mean() == 1.0:
                 del s['noise_mask']
 
-        return io.NodeOutput(output_list)
+        return (output_list,)
 
-class ImageRebatch(io.ComfyNode):
+class ImageRebatch:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="RebatchImages",
-            display_name="Rebatch Images",
-            category="image/batch",
-            is_input_list=True,
-            inputs=[
-                io.Image.Input("images"),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-            ],
-            outputs=[
-                io.Image.Output(is_output_list=True),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "images": ("IMAGE",),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                              }}
+    RETURN_TYPES = ("IMAGE",)
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True, )
 
-    @classmethod
-    def execute(cls, images, batch_size):
+    FUNCTION = "rebatch"
+
+    CATEGORY = "image/batch"
+
+    def rebatch(self, images, batch_size):
         batch_size = batch_size[0]
 
         output_list = []
@@ -137,17 +125,14 @@ class ImageRebatch(io.ComfyNode):
         for i in range(0, len(all_images), batch_size):
             output_list.append(torch.cat(all_images[i:i+batch_size], dim=0))
 
-        return io.NodeOutput(output_list)
+        return (output_list,)
 
+NODE_CLASS_MAPPINGS = {
+    "RebatchLatents": LatentRebatch,
+    "RebatchImages": ImageRebatch,
+}
 
-class RebatchExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            LatentRebatch,
-            ImageRebatch,
-        ]
-
-
-async def comfy_entrypoint() -> RebatchExtension:
-    return RebatchExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "RebatchLatents": "Rebatch Latents",
+    "RebatchImages": "Rebatch Images",
+}

comfy_extras/nodes_sag.py

@@ -2,13 +2,10 @@ import torch
 from torch import einsum
 import torch.nn.functional as F
 import math
-from typing_extensions import override
 
 from einops import rearrange, repeat
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.samplers
-from comfy_api.latest import ComfyExtension, io
-
 
 # from comfy/ldm/modules/attention.py
 # but modified to return attention scores as well as output
@@ -107,26 +104,19 @@ def gaussian_blur_2d(img, kernel_size, sigma):
     img = F.conv2d(img, kernel2d, groups=img.shape[-3])
     return img
 
-class SelfAttentionGuidance(io.ComfyNode):
+class SelfAttentionGuidance:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SelfAttentionGuidance",
-            display_name="Self-Attention Guidance",
-            category="_for_testing",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("scale", default=0.5, min=-2.0, max=5.0, step=0.01),
-                io.Float.Input("blur_sigma", default=2.0, min=0.0, max=10.0, step=0.1),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                             "scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 5.0, "step": 0.01}),
+                             "blur_sigma": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, scale, blur_sigma):
+    CATEGORY = "_for_testing"
+
+    def patch(self, model, scale, blur_sigma):
         m = model.clone()
 
         attn_scores = None
@@ -180,16 +170,12 @@ class SelfAttentionGuidance(io.ComfyNode):
         # unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch
         m.set_model_attn1_replace(attn_and_record, "middle", 0, 0)
 
-        return io.NodeOutput(m)
+        return (m, )
 
+NODE_CLASS_MAPPINGS = {
+    "SelfAttentionGuidance": SelfAttentionGuidance,
+}
 
-class SagExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            SelfAttentionGuidance,
-        ]
-
-
-async def comfy_entrypoint() -> SagExtension:
-    return SagExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SelfAttentionGuidance": "Self-Attention Guidance",
+}

comfy_extras/nodes_sd3.py

@@ -3,83 +3,64 @@ import comfy.sd
 import comfy.model_management
 import nodes
 import torch
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
-from comfy_extras.nodes_slg import SkipLayerGuidanceDiT
+import comfy_extras.nodes_slg
 
 
-class TripleCLIPLoader(io.ComfyNode):
+class TripleCLIPLoader:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="TripleCLIPLoader",
-            category="advanced/loaders",
-            description="[Recipes]\n\nsd3: clip-l, clip-g, t5",
-            inputs=[
-                io.Combo.Input("clip_name1", options=folder_paths.get_filename_list("text_encoders")),
-                io.Combo.Input("clip_name2", options=folder_paths.get_filename_list("text_encoders")),
-                io.Combo.Input("clip_name3", options=folder_paths.get_filename_list("text_encoders")),
-            ],
-            outputs=[
-                io.Clip.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "clip_name1": (folder_paths.get_filename_list("text_encoders"), ), "clip_name2": (folder_paths.get_filename_list("text_encoders"), ), "clip_name3": (folder_paths.get_filename_list("text_encoders"), )
+                             }}
+    RETURN_TYPES = ("CLIP",)
+    FUNCTION = "load_clip"
 
-    @classmethod
-    def execute(cls, clip_name1, clip_name2, clip_name3) -> io.NodeOutput:
+    CATEGORY = "advanced/loaders"
+
+    DESCRIPTION = "[Recipes]\n\nsd3: clip-l, clip-g, t5"
+
+    def load_clip(self, clip_name1, clip_name2, clip_name3):
         clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
         clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
         clip_path3 = folder_paths.get_full_path_or_raise("text_encoders", clip_name3)
         clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2, clip_path3], embedding_directory=folder_paths.get_folder_paths("embeddings"))
-        return io.NodeOutput(clip)
-
-    load_clip = execute  # TODO: remove
+        return (clip,)
 
 
-class EmptySD3LatentImage(io.ComfyNode):
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="EmptySD3LatentImage",
-            category="latent/sd3",
-            inputs=[
-                io.Int.Input("width", default=1024, min=16, max=nodes.MAX_RESOLUTION, step=16),
-                io.Int.Input("height", default=1024, min=16, max=nodes.MAX_RESOLUTION, step=16),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-            ],
-            outputs=[
-                io.Latent.Output(),
-            ],
-        )
+class EmptySD3LatentImage:
+    def __init__(self):
+        self.device = comfy.model_management.intermediate_device()
 
     @classmethod
-    def execute(cls, width, height, batch_size=1) -> io.NodeOutput:
-        latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=comfy.model_management.intermediate_device())
-        return io.NodeOutput({"samples":latent})
+    def INPUT_TYPES(s):
+        return {"required": { "width": ("INT", {"default": 1024, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                              "height": ("INT", {"default": 1024, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "generate"
 
-    generate = execute  # TODO: remove
+    CATEGORY = "latent/sd3"
+
+    def generate(self, width, height, batch_size=1):
+        latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=self.device)
+        return ({"samples":latent}, )
 
 
-class CLIPTextEncodeSD3(io.ComfyNode):
+class CLIPTextEncodeSD3:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="CLIPTextEncodeSD3",
-            category="advanced/conditioning",
-            inputs=[
-                io.Clip.Input("clip"),
-                io.String.Input("clip_l", multiline=True, dynamic_prompts=True),
-                io.String.Input("clip_g", multiline=True, dynamic_prompts=True),
-                io.String.Input("t5xxl", multiline=True, dynamic_prompts=True),
-                io.Combo.Input("empty_padding", options=["none", "empty_prompt"]),
-            ],
-            outputs=[
-                io.Conditioning.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP", ),
+            "clip_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "clip_g": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "t5xxl": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            "empty_padding": (["none", "empty_prompt"], )
+            }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
 
-    @classmethod
-    def execute(cls, clip, clip_l, clip_g, t5xxl, empty_padding) -> io.NodeOutput:
+    CATEGORY = "advanced/conditioning"
+
+    def encode(self, clip, clip_l, clip_g, t5xxl, empty_padding):
         no_padding = empty_padding == "none"
 
         tokens = clip.tokenize(clip_g)
@@ -101,112 +82,57 @@ class CLIPTextEncodeSD3(io.ComfyNode):
                 tokens["l"] += empty["l"]
             while len(tokens["l"]) > len(tokens["g"]):
                 tokens["g"] += empty["g"]
-        return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens))
-
-    encode = execute  # TODO: remove
+        return (clip.encode_from_tokens_scheduled(tokens), )
 
 
-class ControlNetApplySD3(io.ComfyNode):
+class ControlNetApplySD3(nodes.ControlNetApplyAdvanced):
     @classmethod
-    def define_schema(cls) -> io.Schema:
-        return io.Schema(
-            node_id="ControlNetApplySD3",
-            display_name="Apply Controlnet with VAE",
-            category="conditioning/controlnet",
-            inputs=[
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.ControlNet.Input("control_net"),
-                io.Vae.Input("vae"),
-                io.Image.Input("image"),
-                io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01),
-                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-            ],
-            is_deprecated=True,
-        )
-
-    @classmethod
-    def execute(cls, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None) -> io.NodeOutput:
-        if strength == 0:
-            return io.NodeOutput(positive, negative)
-
-        control_hint = image.movedim(-1, 1)
-        cnets = {}
-
-        out = []
-        for conditioning in [positive, negative]:
-            c = []
-            for t in conditioning:
-                d = t[1].copy()
-
-                prev_cnet = d.get('control', None)
-                if prev_cnet in cnets:
-                    c_net = cnets[prev_cnet]
-                else:
-                    c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent),
-                                                             vae=vae, extra_concat=[])
-                    c_net.set_previous_controlnet(prev_cnet)
-                    cnets[prev_cnet] = c_net
-
-                d['control'] = c_net
-                d['control_apply_to_uncond'] = False
-                n = [t[0], d]
-                c.append(n)
-            out.append(c)
-        return io.NodeOutput(out[0], out[1])
-
-    apply_controlnet = execute  # TODO: remove
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "control_net": ("CONTROL_NET", ),
+                             "vae": ("VAE", ),
+                             "image": ("IMAGE", ),
+                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
+                             }}
+    CATEGORY = "conditioning/controlnet"
+    DEPRECATED = True
 
 
-class SkipLayerGuidanceSD3(io.ComfyNode):
+class SkipLayerGuidanceSD3(comfy_extras.nodes_slg.SkipLayerGuidanceDiT):
     '''
     Enhance guidance towards detailed dtructure by having another set of CFG negative with skipped layers.
     Inspired by Perturbed Attention Guidance (https://arxiv.org/abs/2403.17377)
     Experimental implementation by Dango233@StabilityAI.
     '''
-
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SkipLayerGuidanceSD3",
-            category="advanced/guidance",
-            description="Generic version of SkipLayerGuidance node that can be used on every DiT model.",
-            inputs=[
-                io.Model.Input("model"),
-                io.String.Input("layers", default="7, 8, 9", multiline=False),
-                io.Float.Input("scale", default=3.0, min=0.0, max=10.0, step=0.1),
-                io.Float.Input("start_percent", default=0.01, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("end_percent", default=0.15, min=0.0, max=1.0, step=0.001),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL", ),
+                             "layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+                             "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+                             "start_percent": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001})
+                                }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "skip_guidance_sd3"
 
-    @classmethod
-    def execute(cls, model, layers, scale, start_percent, end_percent) -> io.NodeOutput:
-        return SkipLayerGuidanceDiT().execute(model=model, scale=scale, start_percent=start_percent, end_percent=end_percent, double_layers=layers)
+    CATEGORY = "advanced/guidance"
 
-    skip_guidance_sd3 = execute  # TODO: remove
+    def skip_guidance_sd3(self, model, layers, scale, start_percent, end_percent):
+        return self.skip_guidance(model=model, scale=scale, start_percent=start_percent, end_percent=end_percent, double_layers=layers)
 
 
-class SD3Extension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            TripleCLIPLoader,
-            EmptySD3LatentImage,
-            CLIPTextEncodeSD3,
-            ControlNetApplySD3,
-            SkipLayerGuidanceSD3,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "TripleCLIPLoader": TripleCLIPLoader,
+    "EmptySD3LatentImage": EmptySD3LatentImage,
+    "CLIPTextEncodeSD3": CLIPTextEncodeSD3,
+    "ControlNetApplySD3": ControlNetApplySD3,
+    "SkipLayerGuidanceSD3": SkipLayerGuidanceSD3,
+}
 
-
-async def comfy_entrypoint() -> SD3Extension:
-    return SD3Extension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # Sampling
+    "ControlNetApplySD3": "Apply Controlnet with VAE",
+}

comfy_extras/nodes_sdupscale.py

@@ -1,31 +1,23 @@
-from typing_extensions import override
-
 import torch
 import comfy.utils
-from comfy_api.latest import ComfyExtension, io
 
-class SD_4XUpscale_Conditioning(io.ComfyNode):
+class SD_4XUpscale_Conditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SD_4XUpscale_Conditioning",
-            category="conditioning/upscale_diffusion",
-            inputs=[
-                io.Image.Input("images"),
-                io.Conditioning.Input("positive"),
-                io.Conditioning.Input("negative"),
-                io.Float.Input("scale_ratio", default=4.0, min=0.0, max=10.0, step=0.01),
-                io.Float.Input("noise_augmentation", default=0.0, min=0.0, max=1.0, step=0.001),
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent"),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "images": ("IMAGE",),
+                              "positive": ("CONDITIONING",),
+                              "negative": ("CONDITIONING",),
+                              "scale_ratio": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                              "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
 
-    @classmethod
-    def execute(cls, images, positive, negative, scale_ratio, noise_augmentation):
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/upscale_diffusion"
+
+    def encode(self, images, positive, negative, scale_ratio, noise_augmentation):
         width = max(1, round(images.shape[-2] * scale_ratio))
         height = max(1, round(images.shape[-3] * scale_ratio))
 
@@ -47,16 +39,8 @@ class SD_4XUpscale_Conditioning(io.ComfyNode):
             out_cn.append(n)
 
         latent = torch.zeros([images.shape[0], 4, height // 4, width // 4])
-        return io.NodeOutput(out_cp, out_cn, {"samples":latent})
+        return (out_cp, out_cn, {"samples":latent})
 
-
-class SdUpscaleExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            SD_4XUpscale_Conditioning,
-        ]
-
-
-async def comfy_entrypoint() -> SdUpscaleExtension:
-    return SdUpscaleExtension()
+NODE_CLASS_MAPPINGS = {
+    "SD_4XUpscale_Conditioning": SD_4XUpscale_Conditioning,
+}

comfy_extras/nodes_slg.py

@@ -1,40 +1,33 @@
 import comfy.model_patcher
 import comfy.samplers
 import re
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 
-class SkipLayerGuidanceDiT(io.ComfyNode):
+class SkipLayerGuidanceDiT:
     '''
     Enhance guidance towards detailed dtructure by having another set of CFG negative with skipped layers.
     Inspired by Perturbed Attention Guidance (https://arxiv.org/abs/2403.17377)
     Original experimental implementation for SD3 by Dango233@StabilityAI.
     '''
-
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SkipLayerGuidanceDiT",
-            category="advanced/guidance",
-            description="Generic version of SkipLayerGuidance node that can be used on every DiT model.",
-            is_experimental=True,
-            inputs=[
-                io.Model.Input("model"),
-                io.String.Input("double_layers", default="7, 8, 9"),
-                io.String.Input("single_layers", default="7, 8, 9"),
-                io.Float.Input("scale", default=3.0, min=0.0, max=10.0, step=0.1),
-                io.Float.Input("start_percent", default=0.01, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("end_percent", default=0.15, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("rescaling_scale", default=0.0, min=0.0, max=10.0, step=0.01),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL", ),
+                             "double_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+                             "single_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+                             "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+                             "start_percent": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "rescaling_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                                }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "skip_guidance"
+    EXPERIMENTAL = True
 
-    @classmethod
-    def execute(cls, model, scale, start_percent, end_percent, double_layers="", single_layers="", rescaling_scale=0) -> io.NodeOutput:
+    DESCRIPTION = "Generic version of SkipLayerGuidance node that can be used on every DiT model."
+
+    CATEGORY = "advanced/guidance"
+
+    def skip_guidance(self, model, scale, start_percent, end_percent, double_layers="", single_layers="", rescaling_scale=0):
         # check if layer is comma separated integers
         def skip(args, extra_args):
             return args
@@ -50,7 +43,7 @@ class SkipLayerGuidanceDiT(io.ComfyNode):
         single_layers = [int(i) for i in single_layers]
 
         if len(double_layers) == 0 and len(single_layers) == 0:
-            return io.NodeOutput(model)
+            return (model, )
 
         def post_cfg_function(args):
             model = args["model"]
@@ -83,36 +76,29 @@ class SkipLayerGuidanceDiT(io.ComfyNode):
         m = model.clone()
         m.set_model_sampler_post_cfg_function(post_cfg_function)
 
-        return io.NodeOutput(m)
+        return (m, )
 
-    skip_guidance = execute  # TODO: remove
-
-
-class SkipLayerGuidanceDiTSimple(io.ComfyNode):
+class SkipLayerGuidanceDiTSimple:
     '''
     Simple version of the SkipLayerGuidanceDiT node that only modifies the uncond pass.
     '''
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SkipLayerGuidanceDiTSimple",
-            category="advanced/guidance",
-            description="Simple version of the SkipLayerGuidanceDiT node that only modifies the uncond pass.",
-            is_experimental=True,
-            inputs=[
-                io.Model.Input("model"),
-                io.String.Input("double_layers", default="7, 8, 9"),
-                io.String.Input("single_layers", default="7, 8, 9"),
-                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001),
-                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001),
-            ],
-            outputs=[
-                io.Model.Output(),
-            ],
-        )
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL", ),
+                             "double_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+                             "single_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                                }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "skip_guidance"
+    EXPERIMENTAL = True
 
-    @classmethod
-    def execute(cls, model, start_percent, end_percent, double_layers="", single_layers="") -> io.NodeOutput:
+    DESCRIPTION = "Simple version of the SkipLayerGuidanceDiT node that only modifies the uncond pass."
+
+    CATEGORY = "advanced/guidance"
+
+    def skip_guidance(self, model, start_percent, end_percent, double_layers="", single_layers=""):
         def skip(args, extra_args):
             return args
 
@@ -127,7 +113,7 @@ class SkipLayerGuidanceDiTSimple(io.ComfyNode):
         single_layers = [int(i) for i in single_layers]
 
         if len(double_layers) == 0 and len(single_layers) == 0:
-            return io.NodeOutput(model)
+            return (model, )
 
         def calc_cond_batch_function(args):
             x = args["input"]
@@ -158,19 +144,9 @@ class SkipLayerGuidanceDiTSimple(io.ComfyNode):
         m = model.clone()
         m.set_model_sampler_calc_cond_batch_function(calc_cond_batch_function)
 
-        return io.NodeOutput(m)
+        return (m, )
 
-    skip_guidance = execute  # TODO: remove
-
-
-class SkipLayerGuidanceExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            SkipLayerGuidanceDiT,
-            SkipLayerGuidanceDiTSimple,
-        ]
-
-
-async def comfy_entrypoint() -> SkipLayerGuidanceExtension:
-    return SkipLayerGuidanceExtension()
+NODE_CLASS_MAPPINGS = {
+    "SkipLayerGuidanceDiT": SkipLayerGuidanceDiT,
+    "SkipLayerGuidanceDiTSimple": SkipLayerGuidanceDiTSimple,
+}

comfy_extras/nodes_stable3d.py

@@ -1,8 +1,6 @@
 import torch
 import nodes
 import comfy.utils
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 
 def camera_embeddings(elevation, azimuth):
     elevation = torch.as_tensor([elevation])
@@ -22,31 +20,26 @@ def camera_embeddings(elevation, azimuth):
     return embeddings
 
 
-class StableZero123_Conditioning(io.ComfyNode):
+class StableZero123_Conditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="StableZero123_Conditioning",
-            category="conditioning/3d_models",
-            inputs=[
-                io.ClipVision.Input("clip_vision"),
-                io.Image.Input("init_image"),
-                io.Vae.Input("vae"),
-                io.Int.Input("width", default=256, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("height", default=256, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-                io.Float.Input("elevation", default=0.0, min=-180.0, max=180.0, step=0.1, round=False),
-                io.Float.Input("azimuth", default=0.0, min=-180.0, max=180.0, step=0.1, round=False)
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent")
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "clip_vision": ("CLIP_VISION",),
+                              "init_image": ("IMAGE",),
+                              "vae": ("VAE",),
+                              "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
 
-    @classmethod
-    def execute(cls, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth) -> io.NodeOutput:
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/3d_models"
+
+    def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth):
         output = clip_vision.encode_image(init_image)
         pooled = output.image_embeds.unsqueeze(0)
         pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
@@ -58,35 +51,30 @@ class StableZero123_Conditioning(io.ComfyNode):
         positive = [[cond, {"concat_latent_image": t}]]
         negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]]
         latent = torch.zeros([batch_size, 4, height // 8, width // 8])
-        return io.NodeOutput(positive, negative, {"samples":latent})
+        return (positive, negative, {"samples":latent})
 
-class StableZero123_Conditioning_Batched(io.ComfyNode):
+class StableZero123_Conditioning_Batched:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="StableZero123_Conditioning_Batched",
-            category="conditioning/3d_models",
-            inputs=[
-                io.ClipVision.Input("clip_vision"),
-                io.Image.Input("init_image"),
-                io.Vae.Input("vae"),
-                io.Int.Input("width", default=256, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("height", default=256, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("batch_size", default=1, min=1, max=4096),
-                io.Float.Input("elevation", default=0.0, min=-180.0, max=180.0, step=0.1, round=False),
-                io.Float.Input("azimuth", default=0.0, min=-180.0, max=180.0, step=0.1, round=False),
-                io.Float.Input("elevation_batch_increment", default=0.0, min=-180.0, max=180.0, step=0.1, round=False),
-                io.Float.Input("azimuth_batch_increment", default=0.0, min=-180.0, max=180.0, step=0.1, round=False)
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent")
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "clip_vision": ("CLIP_VISION",),
+                              "init_image": ("IMAGE",),
+                              "vae": ("VAE",),
+                              "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
 
-    @classmethod
-    def execute(cls, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth, elevation_batch_increment, azimuth_batch_increment) -> io.NodeOutput:
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/3d_models"
+
+    def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth, elevation_batch_increment, azimuth_batch_increment):
         output = clip_vision.encode_image(init_image)
         pooled = output.image_embeds.unsqueeze(0)
         pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
@@ -105,32 +93,27 @@ class StableZero123_Conditioning_Batched(io.ComfyNode):
         positive = [[cond, {"concat_latent_image": t}]]
         negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]]
         latent = torch.zeros([batch_size, 4, height // 8, width // 8])
-        return io.NodeOutput(positive, negative, {"samples":latent, "batch_index": [0] * batch_size})
+        return (positive, negative, {"samples":latent, "batch_index": [0] * batch_size})
 
-class SV3D_Conditioning(io.ComfyNode):
+class SV3D_Conditioning:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="SV3D_Conditioning",
-            category="conditioning/3d_models",
-            inputs=[
-                io.ClipVision.Input("clip_vision"),
-                io.Image.Input("init_image"),
-                io.Vae.Input("vae"),
-                io.Int.Input("width", default=576, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("height", default=576, min=16, max=nodes.MAX_RESOLUTION, step=8),
-                io.Int.Input("video_frames", default=21, min=1, max=4096),
-                io.Float.Input("elevation", default=0.0, min=-90.0, max=90.0, step=0.1, round=False)
-            ],
-            outputs=[
-                io.Conditioning.Output(display_name="positive"),
-                io.Conditioning.Output(display_name="negative"),
-                io.Latent.Output(display_name="latent")
-            ]
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "clip_vision": ("CLIP_VISION",),
+                              "init_image": ("IMAGE",),
+                              "vae": ("VAE",),
+                              "width": ("INT", {"default": 576, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 576, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
+                              "video_frames": ("INT", {"default": 21, "min": 1, "max": 4096}),
+                              "elevation": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1, "round": False}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
 
-    @classmethod
-    def execute(cls, clip_vision, init_image, vae, width, height, video_frames, elevation) -> io.NodeOutput:
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/3d_models"
+
+    def encode(self, clip_vision, init_image, vae, width, height, video_frames, elevation):
         output = clip_vision.encode_image(init_image)
         pooled = output.image_embeds.unsqueeze(0)
         pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
@@ -150,17 +133,11 @@ class SV3D_Conditioning(io.ComfyNode):
         positive = [[pooled, {"concat_latent_image": t, "elevation": elevations, "azimuth": azimuths}]]
         negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t), "elevation": elevations, "azimuth": azimuths}]]
         latent = torch.zeros([video_frames, 4, height // 8, width // 8])
-        return io.NodeOutput(positive, negative, {"samples":latent})
+        return (positive, negative, {"samples":latent})
 
 
-class Stable3DExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            StableZero123_Conditioning,
-            StableZero123_Conditioning_Batched,
-            SV3D_Conditioning,
-        ]
-
-async def comfy_entrypoint() -> Stable3DExtension:
-    return Stable3DExtension()
+NODE_CLASS_MAPPINGS = {
+    "StableZero123_Conditioning": StableZero123_Conditioning,
+    "StableZero123_Conditioning_Batched": StableZero123_Conditioning_Batched,
+    "SV3D_Conditioning": SV3D_Conditioning,
+}

comfy_extras/nodes_tcfg.py

@@ -1,9 +1,8 @@
 # TCFG: Tangential Damping Classifier-free Guidance - (arXiv: https://arxiv.org/abs/2503.18137)
 
-from typing_extensions import override
 import torch
 
-from comfy_api.latest import ComfyExtension, io
+from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
 
 
 def score_tangential_damping(cond_score: torch.Tensor, uncond_score: torch.Tensor) -> torch.Tensor:
@@ -27,24 +26,23 @@ def score_tangential_damping(cond_score: torch.Tensor, uncond_score: torch.Tenso
     return uncond_score_td.reshape_as(uncond_score).to(uncond_score.dtype)
 
 
-class TCFG(io.ComfyNode):
+class TCFG(ComfyNodeABC):
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="TCFG",
-            display_name="Tangential Damping CFG",
-            category="advanced/guidance",
-            description="TCFG – Tangential Damping CFG (2503.18137)\n\nRefine the uncond (negative) to align with the cond (positive) for improving quality.",
-            inputs=[
-                io.Model.Input("model"),
-            ],
-            outputs=[
-                io.Model.Output(display_name="patched_model"),
-            ],
-        )
+    def INPUT_TYPES(cls) -> InputTypeDict:
+        return {
+            "required": {
+                "model": (IO.MODEL, {}),
+            }
+        }
 
-    @classmethod
-    def execute(cls, model):
+    RETURN_TYPES = (IO.MODEL,)
+    RETURN_NAMES = ("patched_model",)
+    FUNCTION = "patch"
+
+    CATEGORY = "advanced/guidance"
+    DESCRIPTION = "TCFG – Tangential Damping CFG (2503.18137)\n\nRefine the uncond (negative) to align with the cond (positive) for improving quality."
+
+    def patch(self, model):
         m = model.clone()
 
         def tangential_damping_cfg(args):
@@ -61,16 +59,13 @@ class TCFG(io.ComfyNode):
             return [cond_pred, uncond_pred_td] + conds_out[2:]
 
         m.set_model_sampler_pre_cfg_function(tangential_damping_cfg)
-        return io.NodeOutput(m)
+        return (m,)
 
 
-class TcfgExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            TCFG,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "TCFG": TCFG,
+}
 
-
-async def comfy_entrypoint() -> TcfgExtension:
-    return TcfgExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "TCFG": "Tangential Damping CFG",
+}

comfy_extras/nodes_tomesd.py

@@ -1,9 +1,7 @@
 #Taken from: https://github.com/dbolya/tomesd
 
 import torch
-from typing import Tuple, Callable, Optional
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
+from typing import Tuple, Callable
 import math
 
 def do_nothing(x: torch.Tensor, mode:str=None):
@@ -146,45 +144,33 @@ def get_functions(x, ratio, original_shape):
 
 
 
-class TomePatchModel(io.ComfyNode):
+class TomePatchModel:
     @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="TomePatchModel",
-            category="model_patches/unet",
-            inputs=[
-                io.Model.Input("model"),
-                io.Float.Input("ratio", default=0.3, min=0.0, max=1.0, step=0.01),
-            ],
-            outputs=[io.Model.Output()],
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, ratio) -> io.NodeOutput:
-        u: Optional[Callable] = None
+    CATEGORY = "model_patches/unet"
+
+    def patch(self, model, ratio):
+        self.u = None
         def tomesd_m(q, k, v, extra_options):
-            nonlocal u
             #NOTE: In the reference code get_functions takes x (input of the transformer block) as the argument instead of q
             #however from my basic testing it seems that using q instead gives better results
-            m, u = get_functions(q, ratio, extra_options["original_shape"])
+            m, self.u = get_functions(q, ratio, extra_options["original_shape"])
             return m(q), k, v
         def tomesd_u(n, extra_options):
-            nonlocal u
-            return u(n)
+            return self.u(n)
 
         m = model.clone()
         m.set_model_attn1_patch(tomesd_m)
         m.set_model_attn1_output_patch(tomesd_u)
-        return io.NodeOutput(m)
+        return (m, )
 
 
-class TomePatchModelExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            TomePatchModel,
-        ]
-
-
-async def comfy_entrypoint() -> TomePatchModelExtension:
-    return TomePatchModelExtension()
+NODE_CLASS_MAPPINGS = {
+    "TomePatchModel": TomePatchModel,
+}

comfy_extras/nodes_torch_compile.py

@@ -1,39 +1,23 @@
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io
 from comfy_api.torch_helpers import set_torch_compile_wrapper
 
 
-class TorchCompileModel(io.ComfyNode):
+class TorchCompileModel:
     @classmethod
-    def define_schema(cls) -> io.Schema:
-        return io.Schema(
-            node_id="TorchCompileModel",
-            category="_for_testing",
-            inputs=[
-                io.Model.Input("model"),
-                io.Combo.Input(
-                    "backend",
-                    options=["inductor", "cudagraphs"],
-                ),
-            ],
-            outputs=[io.Model.Output()],
-            is_experimental=True,
-        )
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                             "backend": (["inductor", "cudagraphs"],),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
 
-    @classmethod
-    def execute(cls, model, backend) -> io.NodeOutput:
+    CATEGORY = "_for_testing"
+    EXPERIMENTAL = True
+
+    def patch(self, model, backend):
         m = model.clone()
         set_torch_compile_wrapper(model=m, backend=backend)
-        return io.NodeOutput(m)
+        return (m, )
 
-
-class TorchCompileExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [
-            TorchCompileModel,
-        ]
-
-
-async def comfy_entrypoint() -> TorchCompileExtension:
-    return TorchCompileExtension()
+NODE_CLASS_MAPPINGS = {
+    "TorchCompileModel": TorchCompileModel,
+}