diff --git a/.github/workflows/pullrequest-ci-run.yml b/.github/workflows/pullrequest-ci-run.yml
new file mode 100644
index 000000000..fd1ba16e2
--- /dev/null
+++ b/.github/workflows/pullrequest-ci-run.yml
@@ -0,0 +1,37 @@
+# This is the GitHub Workflow that drives full-GPU-enabled tests of pull requests to ComfyUI, when the 'Run-CI-Test' label is added
+# Results are reported as checkmarks on the commits, as well as onto https://ci.comfy.org/
+name: Pull Request CI Workflow Runs
+on:
+ pull_request_target:
+ types: [labeled]
+
+jobs:
+ pr-test-stable:
+ if: ${{ github.event.label.name == 'Run-CI-Test' }}
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [macos, linux, windows]
+ python_version: ["3.9", "3.10", "3.11", "3.12"]
+ cuda_version: ["12.1"]
+ torch_version: ["stable"]
+ include:
+ - os: macos
+ runner_label: [self-hosted, macOS]
+ flags: "--use-pytorch-cross-attention"
+ - os: linux
+ runner_label: [self-hosted, Linux]
+ flags: ""
+ - os: windows
+ runner_label: [self-hosted, win]
+ flags: ""
+ runs-on: ${{ matrix.runner_label }}
+ steps:
+ - name: Test Workflows
+ uses: comfy-org/comfy-action@main
+ with:
+ os: ${{ matrix.os }}
+ python_version: ${{ matrix.python_version }}
+ torch_version: ${{ matrix.torch_version }}
+ google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+ comfyui_flags: ${{ matrix.flags }}
diff --git a/.github/workflows/stable-release.yml b/.github/workflows/stable-release.yml
index 19035c02c..658816afe 100644
--- a/.github/workflows/stable-release.yml
+++ b/.github/workflows/stable-release.yml
@@ -2,9 +2,28 @@
name: "Release Stable Version"
on:
- push:
- tags:
- - 'v*'
+ workflow_dispatch:
+ inputs:
+ git_tag:
+ description: 'Git tag'
+ required: true
+ type: string
+ cu:
+ description: 'CUDA version'
+ required: true
+ type: string
+ default: "121"
+ python_minor:
+ description: 'Python minor version'
+ required: true
+ type: string
+ default: "11"
+ python_patch:
+ description: 'Python patch version'
+ required: true
+ type: string
+ default: "9"
+
jobs:
package_comfy_windows:
@@ -13,69 +32,44 @@ jobs:
packages: "write"
pull-requests: "read"
runs-on: windows-latest
- strategy:
- matrix:
- python_version: [3.11.8]
- cuda_version: [121]
steps:
- - name: Calculate Minor Version
- shell: bash
- run: |
- # Extract the minor version from the Python version
- MINOR_VERSION=$(echo "${{ matrix.python_version }}" | cut -d'.' -f2)
- echo "MINOR_VERSION=$MINOR_VERSION" >> $GITHUB_ENV
- - name: Setup Python
- uses: actions/setup-python@v5
- with:
- python-version: ${{ matrix.python_version }}
-
- uses: actions/checkout@v4
with:
+ ref: ${{ inputs.git_tag }}
fetch-depth: 0
persist-credentials: false
+ - uses: actions/cache/restore@v4
+ id: cache
+ with:
+ path: |
+ cu${{ inputs.cu }}_python_deps.tar
+ update_comfyui_and_python_dependencies.bat
+ key: ${{ runner.os }}-build-cu${{ inputs.cu }}-${{ inputs.python_minor }}
- 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 torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu${{ matrix.cuda_version }} -r ../ComfyUI/requirements.txt pygit2
- pause" > update_comfyui_and_python_dependencies.bat
-
- python -m pip wheel --no-cache-dir torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu${{ matrix.cuda_version }} -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
- mv temp_wheel_dir cu${{ matrix.cuda_version }}_python_deps
- mv cu${{ matrix.cuda_version }}_python_deps ../
+ mv cu${{ inputs.cu }}_python_deps.tar ../
mv update_comfyui_and_python_dependencies.bat ../
cd ..
+ tar xf cu${{ inputs.cu }}_python_deps.tar
pwd
ls
-
+
+ - shell: bash
+ run: |
+ cd ..
cp -r ComfyUI ComfyUI_copy
- curl https://www.python.org/ftp/python/${{ matrix.python_version }}/python-${{ matrix.python_version }}-embed-amd64.zip -o python_embeded.zip
+ curl https://www.python.org/ftp/python/3.${{ inputs.python_minor }}.${{ inputs.python_patch }}/python-3.${{ inputs.python_minor }}.${{ inputs.python_patch }}-embed-amd64.zip -o python_embeded.zip
unzip python_embeded.zip -d python_embeded
cd python_embeded
echo ${{ env.MINOR_VERSION }}
- echo 'import site' >> ./python3${{ env.MINOR_VERSION }}._pth
+ 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 --version
- echo "Pip version:"
- ./python.exe -m pip --version
+ ./python.exe -s -m pip install ../cu${{ inputs.cu }}_python_deps/*
+ sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
+ cd ..
- set PATH=$PWD/Scripts:$PATH
- echo $PATH
- ./python.exe -s -m pip install ../cu${{ matrix.cuda_version }}_python_deps/*
- sed -i '1i../ComfyUI' ./python3${{ env.MINOR_VERSION }}._pth
- cd ..
-
- git clone https://github.com/comfyanonymous/taesd
+ git clone --depth 1 https://github.com/comfyanonymous/taesd
cp taesd/*.pth ./ComfyUI_copy/models/vae_approx/
mkdir ComfyUI_windows_portable
@@ -104,7 +98,7 @@ jobs:
with:
repo_token: ${{ secrets.GITHUB_TOKEN }}
file: ComfyUI_windows_portable_nvidia.7z
- tag: ${{ github.ref }}
+ tag: ${{ inputs.git_tag }}
overwrite: true
prerelease: true
make_latest: false
diff --git a/.github/workflows/test-browser.yml b/.github/workflows/test-browser.yml
index 7beb0c696..ce0bc37a3 100644
--- a/.github/workflows/test-browser.yml
+++ b/.github/workflows/test-browser.yml
@@ -32,7 +32,7 @@ jobs:
node-version: lts/*
- uses: actions/setup-python@v4
with:
- python-version: '3.10'
+ python-version: '3.8'
- name: Install requirements
run: |
python -m pip install --upgrade pip
diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
new file mode 100644
index 000000000..be7622907
--- /dev/null
+++ b/.github/workflows/test-ci.yml
@@ -0,0 +1,95 @@
+# This is the GitHub Workflow that drives automatic full-GPU-enabled tests of all new commits to the master branch of ComfyUI
+# Results are reported as checkmarks on the commits, as well as onto https://ci.comfy.org/
+name: Full Comfy CI Workflow Runs
+on:
+ push:
+ branches:
+ - master
+ paths-ignore:
+ - 'app/**'
+ - 'input/**'
+ - 'output/**'
+ - 'notebooks/**'
+ - 'script_examples/**'
+ - '.github/**'
+ - 'web/**'
+ workflow_dispatch:
+
+jobs:
+ test-stable:
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [macos, linux, windows]
+ python_version: ["3.9", "3.10", "3.11", "3.12"]
+ cuda_version: ["12.1"]
+ torch_version: ["stable"]
+ include:
+ - os: macos
+ runner_label: [self-hosted, macOS]
+ flags: "--use-pytorch-cross-attention"
+ - os: linux
+ runner_label: [self-hosted, Linux]
+ flags: ""
+ - os: windows
+ runner_label: [self-hosted, win]
+ flags: ""
+ runs-on: ${{ matrix.runner_label }}
+ steps:
+ - name: Test Workflows
+ uses: comfy-org/comfy-action@main
+ with:
+ os: ${{ matrix.os }}
+ python_version: ${{ matrix.python_version }}
+ torch_version: ${{ matrix.torch_version }}
+ google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+ comfyui_flags: ${{ matrix.flags }}
+
+ test-win-nightly:
+ strategy:
+ fail-fast: true
+ matrix:
+ os: [windows]
+ python_version: ["3.9", "3.10", "3.11", "3.12"]
+ cuda_version: ["12.1"]
+ torch_version: ["nightly"]
+ include:
+ - os: windows
+ runner_label: [self-hosted, win]
+ flags: ""
+ runs-on: ${{ matrix.runner_label }}
+ steps:
+ - name: Test Workflows
+ uses: comfy-org/comfy-action@main
+ with:
+ os: ${{ matrix.os }}
+ python_version: ${{ matrix.python_version }}
+ torch_version: ${{ matrix.torch_version }}
+ google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+ comfyui_flags: ${{ matrix.flags }}
+
+ test-unix-nightly:
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [macos, linux]
+ python_version: ["3.11"]
+ cuda_version: ["12.1"]
+ torch_version: ["nightly"]
+ include:
+ - os: macos
+ runner_label: [self-hosted, macOS]
+ flags: "--use-pytorch-cross-attention"
+ - os: linux
+ runner_label: [self-hosted, Linux]
+ flags: ""
+ runs-on: ${{ matrix.runner_label }}
+ steps:
+ - name: Test Workflows
+ uses: comfy-org/comfy-action@main
+ with:
+ os: ${{ matrix.os }}
+ python_version: ${{ matrix.python_version }}
+ torch_version: ${{ matrix.torch_version }}
+ google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+ comfyui_flags: ${{ matrix.flags }}
diff --git a/.github/workflows/windows_release_nightly_pytorch.yml b/.github/workflows/windows_release_nightly_pytorch.yml
index ba388cd44..0b29b4d91 100644
--- a/.github/workflows/windows_release_nightly_pytorch.yml
+++ b/.github/workflows/windows_release_nightly_pytorch.yml
@@ -55,7 +55,7 @@ jobs:
sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
cd ..
- git clone https://github.com/comfyanonymous/taesd
+ git clone --depth 1 https://github.com/comfyanonymous/taesd
cp taesd/*.pth ./ComfyUI_copy/models/vae_approx/
mkdir ComfyUI_windows_portable_nightly_pytorch
diff --git a/.github/workflows/windows_release_package.yml b/.github/workflows/windows_release_package.yml
index 5aed73e55..84d99b8a6 100644
--- a/.github/workflows/windows_release_package.yml
+++ b/.github/workflows/windows_release_package.yml
@@ -66,7 +66,7 @@ jobs:
sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
cd ..
- git clone https://github.com/comfyanonymous/taesd
+ git clone --depth 1 https://github.com/comfyanonymous/taesd
cp taesd/*.pth ./ComfyUI_copy/models/vae_approx/
mkdir ComfyUI_windows_portable
diff --git a/comfy/controlnet.py b/comfy/controlnet.py
index 12e5f16c8..354b00367 100644
--- a/comfy/controlnet.py
+++ b/comfy/controlnet.py
@@ -1,4 +1,24 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
+
import torch
+from enum import Enum
import math
import os
import logging
@@ -13,7 +33,7 @@ import comfy.cldm.cldm
import comfy.t2i_adapter.adapter
import comfy.ldm.cascade.controlnet
import comfy.cldm.mmdit
-
+import comfy.ldm.hydit.controlnet
def broadcast_image_to(tensor, target_batch_size, batched_number):
current_batch_size = tensor.shape[0]
@@ -33,6 +53,10 @@ def broadcast_image_to(tensor, target_batch_size, batched_number):
else:
return torch.cat([tensor] * batched_number, dim=0)
+class StrengthType(Enum):
+ CONSTANT = 1
+ LINEAR_UP = 2
+
class ControlBase:
def __init__(self, device=None):
self.cond_hint_original = None
@@ -51,6 +75,8 @@ class ControlBase:
device = comfy.model_management.get_torch_device()
self.device = device
self.previous_controlnet = None
+ self.extra_conds = []
+ self.strength_type = StrengthType.CONSTANT
def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None):
self.cond_hint_original = cond_hint
@@ -93,6 +119,8 @@ class ControlBase:
c.latent_format = self.latent_format
c.extra_args = self.extra_args.copy()
c.vae = self.vae
+ c.extra_conds = self.extra_conds.copy()
+ c.strength_type = self.strength_type
def inference_memory_requirements(self, dtype):
if self.previous_controlnet is not None:
@@ -113,7 +141,10 @@ class ControlBase:
if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once
applied_to.add(x)
- x *= self.strength
+ if self.strength_type == StrengthType.CONSTANT:
+ x *= self.strength
+ elif self.strength_type == StrengthType.LINEAR_UP:
+ x *= (self.strength ** float(len(control_output) - i))
if x.dtype != output_dtype:
x = x.to(output_dtype)
@@ -142,7 +173,7 @@ class ControlBase:
class ControlNet(ControlBase):
- def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None):
+ def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT):
super().__init__(device)
self.control_model = control_model
self.load_device = load_device
@@ -154,6 +185,8 @@ class ControlNet(ControlBase):
self.model_sampling_current = None
self.manual_cast_dtype = manual_cast_dtype
self.latent_format = latent_format
+ self.extra_conds += extra_conds
+ self.strength_type = strength_type
def get_control(self, x_noisy, t, cond, batched_number):
control_prev = None
@@ -191,13 +224,16 @@ class ControlNet(ControlBase):
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
context = cond.get('crossattn_controlnet', cond['c_crossattn'])
- y = cond.get('y', None)
- if y is not None:
- y = y.to(dtype)
+ extra = self.extra_args.copy()
+ for c in self.extra_conds:
+ temp = cond.get(c, None)
+ if temp is not None:
+ extra[c] = temp.to(dtype)
+
timestep = self.model_sampling_current.timestep(t)
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
- control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, **self.extra_args)
+ control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
return self.control_merge(control, control_prev, output_dtype)
def copy(self):
@@ -286,6 +322,7 @@ class ControlLora(ControlNet):
ControlBase.__init__(self, device)
self.control_weights = control_weights
self.global_average_pooling = global_average_pooling
+ self.extra_conds += ["y"]
def pre_run(self, model, percent_to_timestep_function):
super().pre_run(model, percent_to_timestep_function)
@@ -338,12 +375,8 @@ class ControlLora(ControlNet):
def inference_memory_requirements(self, dtype):
return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
-def load_controlnet_mmdit(sd):
- new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
- model_config = comfy.model_detection.model_config_from_unet(new_sd, "", True)
- num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
- for k in sd:
- new_sd[k] = sd[k]
+def controlnet_config(sd):
+ model_config = comfy.model_detection.model_config_from_unet(sd, "", True)
supported_inference_dtypes = model_config.supported_inference_dtypes
@@ -356,23 +389,49 @@ def load_controlnet_mmdit(sd):
else:
operations = comfy.ops.disable_weight_init
- control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=load_device, dtype=unet_dtype, **controlnet_config)
- missing, unexpected = control_model.load_state_dict(new_sd, strict=False)
+ return model_config, operations, load_device, unet_dtype, manual_cast_dtype
+
+def controlnet_load_state_dict(control_model, sd):
+ missing, unexpected = control_model.load_state_dict(sd, strict=False)
if len(missing) > 0:
logging.warning("missing controlnet keys: {}".format(missing))
if len(unexpected) > 0:
logging.debug("unexpected controlnet keys: {}".format(unexpected))
+ return control_model
+
+def load_controlnet_mmdit(sd):
+ new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
+ model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(new_sd)
+ num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
+ for k in sd:
+ new_sd[k] = sd[k]
+
+ control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=load_device, dtype=unet_dtype, **model_config.unet_config)
+ control_model = controlnet_load_state_dict(control_model, new_sd)
latent_format = comfy.latent_formats.SD3()
latent_format.shift_factor = 0 #SD3 controlnet weirdness
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
return control
+def load_controlnet_hunyuandit(controlnet_data):
+ model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(controlnet_data)
+
+ control_model = comfy.ldm.hydit.controlnet.HunYuanControlNet(operations=operations, device=load_device, dtype=unet_dtype)
+ control_model = controlnet_load_state_dict(control_model, controlnet_data)
+
+ latent_format = comfy.latent_formats.SDXL()
+ extra_conds = ['text_embedding_mask', 'encoder_hidden_states_t5', 'text_embedding_mask_t5', 'image_meta_size', 'style', 'cos_cis_img', 'sin_cis_img']
+ control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)
+ return control
def load_controlnet(ckpt_path, model=None):
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
+ if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT
+ return load_controlnet_hunyuandit(controlnet_data)
+
if "lora_controlnet" in controlnet_data:
return ControlLora(controlnet_data)
diff --git a/comfy/ldm/flux/layers.py b/comfy/ldm/flux/layers.py
index 99f498106..a7b0c55fa 100644
--- a/comfy/ldm/flux/layers.py
+++ b/comfy/ldm/flux/layers.py
@@ -2,12 +2,12 @@ import math
from dataclasses import dataclass
import torch
-from einops import rearrange
from torch import Tensor, nn
from .math import attention, rope
import comfy.ops
+
class EmbedND(nn.Module):
def __init__(self, dim: int, theta: int, axes_dim: list):
super().__init__()
@@ -36,9 +36,7 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
"""
t = time_factor * t
half = dim // 2
- freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
- t.device
- )
+ freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half)
args = t[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
@@ -48,7 +46,6 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
embedding = embedding.to(t)
return embedding
-
class MLPEmbedder(nn.Module):
def __init__(self, in_dim: int, hidden_dim: int, dtype=None, device=None, operations=None):
super().__init__()
@@ -94,14 +91,6 @@ class SelfAttention(nn.Module):
self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
- def forward(self, x: Tensor, pe: Tensor) -> Tensor:
- qkv = self.qkv(x)
- q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
- q, k = self.norm(q, k, v)
- x = attention(q, k, v, pe=pe)
- x = self.proj(x)
- return x
-
@dataclass
class ModulationOut:
@@ -163,31 +152,34 @@ class DoubleStreamBlock(nn.Module):
img_modulated = self.img_norm1(img)
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
img_qkv = self.img_attn.qkv(img_modulated)
- img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ img_q, img_k, img_v = img_qkv.view(img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
# prepare txt for attention
txt_modulated = self.txt_norm1(txt)
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
txt_qkv = self.txt_attn.qkv(txt_modulated)
- txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ txt_q, txt_k, txt_v = txt_qkv.view(txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
# run actual attention
- q = torch.cat((txt_q, img_q), dim=2)
- k = torch.cat((txt_k, img_k), dim=2)
- v = torch.cat((txt_v, img_v), dim=2)
+ attn = attention(torch.cat((txt_q, img_q), dim=2),
+ torch.cat((txt_k, img_k), dim=2),
+ torch.cat((txt_v, img_v), dim=2), pe=pe)
- attn = attention(q, k, v, pe=pe)
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
# calculate the img bloks
- img = img + img_mod1.gate * self.img_attn.proj(img_attn)
- img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+ img += img_mod1.gate * self.img_attn.proj(img_attn)
+ img += img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
# calculate the txt bloks
- txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
- txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+ txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
+ txt += txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+
+ if txt.dtype == torch.float16:
+ txt = txt.clip(-65504, 65504)
+
return img, txt
@@ -232,14 +224,17 @@ class SingleStreamBlock(nn.Module):
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
- q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
q, k = self.norm(q, k, v)
# compute attention
attn = attention(q, k, v, pe=pe)
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
- return x + mod.gate * output
+ x += mod.gate * output
+ if x.dtype == torch.float16:
+ x = x.clip(-65504, 65504)
+ return x
class LastLayer(nn.Module):
diff --git a/comfy/ldm/hydit/attn_layers.py b/comfy/ldm/hydit/attn_layers.py
index 920b84286..e2801f714 100644
--- a/comfy/ldm/hydit/attn_layers.py
+++ b/comfy/ldm/hydit/attn_layers.py
@@ -47,7 +47,7 @@ def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x
def rotate_half(x):
- x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
+ x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
@@ -78,10 +78,9 @@ def apply_rotary_emb(
xk_out = None
if isinstance(freqs_cis, tuple):
cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D]
- cos, sin = cos.to(xq.device), sin.to(xq.device)
- xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+ xq_out = (xq * cos + rotate_half(xq) * sin)
if xk is not None:
- xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+ xk_out = (xk * cos + rotate_half(xk) * sin)
else:
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2]
freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2]
diff --git a/comfy/ldm/hydit/controlnet.py b/comfy/ldm/hydit/controlnet.py
new file mode 100644
index 000000000..cd71fca31
--- /dev/null
+++ b/comfy/ldm/hydit/controlnet.py
@@ -0,0 +1,321 @@
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.utils import checkpoint
+
+from comfy.ldm.modules.diffusionmodules.mmdit import (
+ Mlp,
+ TimestepEmbedder,
+ PatchEmbed,
+ RMSNorm,
+)
+from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
+from .poolers import AttentionPool
+
+import comfy.latent_formats
+from .models import HunYuanDiTBlock, calc_rope
+
+from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
+
+
+class HunYuanControlNet(nn.Module):
+ """
+ HunYuanDiT: Diffusion model with a Transformer backbone.
+
+ Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers.
+
+ Inherit PeftAdapterMixin to be compatible with the PEFT training pipeline.
+
+ Parameters
+ ----------
+ args: argparse.Namespace
+ The arguments parsed by argparse.
+ input_size: tuple
+ The size of the input image.
+ patch_size: int
+ The size of the patch.
+ in_channels: int
+ The number of input channels.
+ hidden_size: int
+ The hidden size of the transformer backbone.
+ depth: int
+ The number of transformer blocks.
+ num_heads: int
+ The number of attention heads.
+ mlp_ratio: float
+ The ratio of the hidden size of the MLP in the transformer block.
+ log_fn: callable
+ The logging function.
+ """
+
+ def __init__(
+ self,
+ input_size: tuple = 128,
+ patch_size: int = 2,
+ in_channels: int = 4,
+ hidden_size: int = 1408,
+ depth: int = 40,
+ num_heads: int = 16,
+ mlp_ratio: float = 4.3637,
+ text_states_dim=1024,
+ text_states_dim_t5=2048,
+ text_len=77,
+ text_len_t5=256,
+ qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details.
+ size_cond=False,
+ use_style_cond=False,
+ learn_sigma=True,
+ norm="layer",
+ log_fn: callable = print,
+ attn_precision=None,
+ dtype=None,
+ device=None,
+ operations=None,
+ **kwargs,
+ ):
+ super().__init__()
+ self.log_fn = log_fn
+ self.depth = depth
+ self.learn_sigma = learn_sigma
+ self.in_channels = in_channels
+ self.out_channels = in_channels * 2 if learn_sigma else in_channels
+ self.patch_size = patch_size
+ self.num_heads = num_heads
+ self.hidden_size = hidden_size
+ self.text_states_dim = text_states_dim
+ self.text_states_dim_t5 = text_states_dim_t5
+ self.text_len = text_len
+ self.text_len_t5 = text_len_t5
+ self.size_cond = size_cond
+ self.use_style_cond = use_style_cond
+ self.norm = norm
+ self.dtype = dtype
+ self.latent_format = comfy.latent_formats.SDXL
+
+ self.mlp_t5 = nn.Sequential(
+ nn.Linear(
+ self.text_states_dim_t5,
+ self.text_states_dim_t5 * 4,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ nn.SiLU(),
+ nn.Linear(
+ self.text_states_dim_t5 * 4,
+ self.text_states_dim,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ )
+ # learnable replace
+ self.text_embedding_padding = nn.Parameter(
+ torch.randn(
+ self.text_len + self.text_len_t5,
+ self.text_states_dim,
+ dtype=dtype,
+ device=device,
+ )
+ )
+
+ # Attention pooling
+ pooler_out_dim = 1024
+ self.pooler = AttentionPool(
+ self.text_len_t5,
+ self.text_states_dim_t5,
+ num_heads=8,
+ output_dim=pooler_out_dim,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+
+ # Dimension of the extra input vectors
+ self.extra_in_dim = pooler_out_dim
+
+ if self.size_cond:
+ # Image size and crop size conditions
+ self.extra_in_dim += 6 * 256
+
+ if self.use_style_cond:
+ # Here we use a default learned embedder layer for future extension.
+ self.style_embedder = nn.Embedding(
+ 1, hidden_size, dtype=dtype, device=device
+ )
+ self.extra_in_dim += hidden_size
+
+ # Text embedding for `add`
+ self.x_embedder = PatchEmbed(
+ input_size,
+ patch_size,
+ in_channels,
+ hidden_size,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ self.t_embedder = TimestepEmbedder(
+ hidden_size, dtype=dtype, device=device, operations=operations
+ )
+ self.extra_embedder = nn.Sequential(
+ operations.Linear(
+ self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device
+ ),
+ nn.SiLU(),
+ operations.Linear(
+ hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device
+ ),
+ )
+
+ # Image embedding
+ num_patches = self.x_embedder.num_patches
+
+ # HUnYuanDiT Blocks
+ self.blocks = nn.ModuleList(
+ [
+ HunYuanDiTBlock(
+ hidden_size=hidden_size,
+ c_emb_size=hidden_size,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ text_states_dim=self.text_states_dim,
+ qk_norm=qk_norm,
+ norm_type=self.norm,
+ skip=False,
+ attn_precision=attn_precision,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ for _ in range(19)
+ ]
+ )
+
+ # Input zero linear for the first block
+ self.before_proj = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
+
+
+ # Output zero linear for the every block
+ self.after_proj_list = nn.ModuleList(
+ [
+
+ operations.Linear(
+ self.hidden_size, self.hidden_size, dtype=dtype, device=device
+ )
+ for _ in range(len(self.blocks))
+ ]
+ )
+
+ def forward(
+ self,
+ x,
+ hint,
+ timesteps,
+ context,#encoder_hidden_states=None,
+ text_embedding_mask=None,
+ encoder_hidden_states_t5=None,
+ text_embedding_mask_t5=None,
+ image_meta_size=None,
+ style=None,
+ return_dict=False,
+ **kwarg,
+ ):
+ """
+ Forward pass of the encoder.
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ (B, D, H, W)
+ t: torch.Tensor
+ (B)
+ encoder_hidden_states: torch.Tensor
+ CLIP text embedding, (B, L_clip, D)
+ text_embedding_mask: torch.Tensor
+ CLIP text embedding mask, (B, L_clip)
+ encoder_hidden_states_t5: torch.Tensor
+ T5 text embedding, (B, L_t5, D)
+ text_embedding_mask_t5: torch.Tensor
+ T5 text embedding mask, (B, L_t5)
+ image_meta_size: torch.Tensor
+ (B, 6)
+ style: torch.Tensor
+ (B)
+ cos_cis_img: torch.Tensor
+ sin_cis_img: torch.Tensor
+ return_dict: bool
+ Whether to return a dictionary.
+ """
+ condition = hint
+ if condition.shape[0] == 1:
+ condition = torch.repeat_interleave(condition, x.shape[0], dim=0)
+
+ text_states = context # 2,77,1024
+ text_states_t5 = encoder_hidden_states_t5 # 2,256,2048
+ text_states_mask = text_embedding_mask.bool() # 2,77
+ text_states_t5_mask = text_embedding_mask_t5.bool() # 2,256
+ b_t5, l_t5, c_t5 = text_states_t5.shape
+ text_states_t5 = self.mlp_t5(text_states_t5.view(-1, c_t5)).view(b_t5, l_t5, -1)
+
+ padding = comfy.ops.cast_to_input(self.text_embedding_padding, text_states)
+
+ text_states[:, -self.text_len :] = torch.where(
+ text_states_mask[:, -self.text_len :].unsqueeze(2),
+ text_states[:, -self.text_len :],
+ padding[: self.text_len],
+ )
+ text_states_t5[:, -self.text_len_t5 :] = torch.where(
+ text_states_t5_mask[:, -self.text_len_t5 :].unsqueeze(2),
+ text_states_t5[:, -self.text_len_t5 :],
+ padding[self.text_len :],
+ )
+
+ text_states = torch.cat([text_states, text_states_t5], dim=1) # 2,205，1024
+
+ # _, _, oh, ow = x.shape
+ # th, tw = oh // self.patch_size, ow // self.patch_size
+
+ # Get image RoPE embedding according to `reso`lution.
+ freqs_cis_img = calc_rope(
+ x, self.patch_size, self.hidden_size // self.num_heads
+ ) # (cos_cis_img, sin_cis_img)
+
+ # ========================= Build time and image embedding =========================
+ t = self.t_embedder(timesteps, dtype=self.dtype)
+ x = self.x_embedder(x)
+
+ # ========================= Concatenate all extra vectors =========================
+ # Build text tokens with pooling
+ extra_vec = self.pooler(encoder_hidden_states_t5)
+
+ # Build image meta size tokens if applicable
+ # if image_meta_size is not None:
+ # image_meta_size = timestep_embedding(image_meta_size.view(-1), 256) # [B * 6, 256]
+ # if image_meta_size.dtype != self.dtype:
+ # image_meta_size = image_meta_size.half()
+ # image_meta_size = image_meta_size.view(-1, 6 * 256)
+ # extra_vec = torch.cat([extra_vec, image_meta_size], dim=1) # [B, D + 6 * 256]
+
+ # Build style tokens
+ if style is not None:
+ style_embedding = self.style_embedder(style)
+ extra_vec = torch.cat([extra_vec, style_embedding], dim=1)
+
+ # Concatenate all extra vectors
+ c = t + self.extra_embedder(extra_vec) # [B, D]
+
+ # ========================= Deal with Condition =========================
+ condition = self.x_embedder(condition)
+
+ # ========================= Forward pass through HunYuanDiT blocks =========================
+ controls = []
+ x = x + self.before_proj(condition) # add condition
+ for layer, block in enumerate(self.blocks):
+ x = block(x, c, text_states, freqs_cis_img)
+ controls.append(self.after_proj_list[layer](x)) # zero linear for output
+
+ return {"output": controls}
diff --git a/comfy/ldm/hydit/models.py b/comfy/ldm/hydit/models.py
index c70dbf92a..f3afaad34 100644
--- a/comfy/ldm/hydit/models.py
+++ b/comfy/ldm/hydit/models.py
@@ -21,6 +21,7 @@ def calc_rope(x, patch_size, head_size):
sub_args = [start, stop, (th, tw)]
# head_size = HUNYUAN_DIT_CONFIG['DiT-g/2']['hidden_size'] // HUNYUAN_DIT_CONFIG['DiT-g/2']['num_heads']
rope = get_2d_rotary_pos_embed(head_size, *sub_args)
+ rope = (rope[0].to(x), rope[1].to(x))
return rope
@@ -91,6 +92,8 @@ class HunYuanDiTBlock(nn.Module):
# Long Skip Connection
if self.skip_linear is not None:
cat = torch.cat([x, skip], dim=-1)
+ if cat.dtype != x.dtype:
+ cat = cat.to(x.dtype)
cat = self.skip_norm(cat)
x = self.skip_linear(cat)
@@ -362,6 +365,8 @@ class HunYuanDiT(nn.Module):
c = t + self.extra_embedder(extra_vec) # [B, D]
controls = None
+ if control:
+ controls = control.get("output", None)
# ========================= Forward pass through HunYuanDiT blocks =========================
skips = []
for layer, block in enumerate(self.blocks):
diff --git a/comfy/lora.py b/comfy/lora.py
index 04e8861c9..3b8b6c162 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -1,3 +1,21 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
import comfy.utils
import logging
@@ -218,11 +236,17 @@ def model_lora_keys_clip(model, key_map={}):
lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config
key_map[lora_key] = k
- for k in sdk: #OneTrainer SD3 lora
- if k.startswith("t5xxl.transformer.") and k.endswith(".weight"):
- l_key = k[len("t5xxl.transformer."):-len(".weight")]
- lora_key = "lora_te3_{}".format(l_key.replace(".", "_"))
- key_map[lora_key] = k
+ for k in sdk:
+ if k.endswith(".weight"):
+ if k.startswith("t5xxl.transformer."):#OneTrainer SD3 lora
+ l_key = k[len("t5xxl.transformer."):-len(".weight")]
+ lora_key = "lora_te3_{}".format(l_key.replace(".", "_"))
+ key_map[lora_key] = k
+ elif k.startswith("hydit_clip.transformer.bert."): #HunyuanDiT Lora
+ l_key = k[len("hydit_clip.transformer.bert."):-len(".weight")]
+ lora_key = "lora_te1_{}".format(l_key.replace(".", "_"))
+ key_map[lora_key] = k
+
k = "clip_g.transformer.text_projection.weight"
if k in sdk:
@@ -245,6 +269,7 @@ def model_lora_keys_unet(model, key_map={}):
key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
key_map["lora_unet_{}".format(key_lora)] = k
key_map["lora_prior_unet_{}".format(key_lora)] = k #cascade lora: TODO put lora key prefix in the model config
+ key_map["{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names
diffusers_keys = comfy.utils.unet_to_diffusers(model.model_config.unet_config)
for k in diffusers_keys:
diff --git a/comfy/model_base.py b/comfy/model_base.py
index d19f5697a..cb6949649 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -74,6 +74,7 @@ class BaseModel(torch.nn.Module):
self.latent_format = model_config.latent_format
self.model_config = model_config
self.manual_cast_dtype = model_config.manual_cast_dtype
+ self.device = device
if not unet_config.get("disable_unet_model_creation", False):
if self.manual_cast_dtype is not None:
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index c47119686..15e6b735c 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -137,8 +137,8 @@ def detect_unet_config(state_dict, key_prefix):
dit_config["hidden_size"] = 3072
dit_config["mlp_ratio"] = 4.0
dit_config["num_heads"] = 24
- dit_config["depth"] = 19
- dit_config["depth_single_blocks"] = 38
+ dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
+ dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
dit_config["axes_dim"] = [16, 56, 56]
dit_config["theta"] = 10000
dit_config["qkv_bias"] = True
diff --git a/comfy/model_management.py b/comfy/model_management.py
index d23501201..b901a444a 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -1,3 +1,21 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
import psutil
import logging
from enum import Enum
@@ -286,8 +304,11 @@ class LoadedModel:
def model_memory(self):
return self.model.model_size()
+ def model_offloaded_memory(self):
+ return self.model.model_size() - self.model.loaded_size()
+
def model_memory_required(self, device):
- if device == self.model.current_device:
+ if device == self.model.current_loaded_device():
return 0
else:
return self.model_memory()
@@ -317,19 +338,41 @@ class LoadedModel:
return self.real_model
def should_reload_model(self, force_patch_weights=False):
- if force_patch_weights and self.model.lowvram_patch_counter > 0:
+ if force_patch_weights and self.model.lowvram_patch_counter() > 0:
return True
return False
- def model_unload(self, unpatch_weights=True):
+ def model_unload(self, memory_to_free=None, unpatch_weights=True):
+ if memory_to_free is not None:
+ if memory_to_free < self.model.loaded_size():
+ self.model.partially_unload(self.model.offload_device, memory_to_free)
+ return False
self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
self.model.model_patches_to(self.model.offload_device)
self.weights_loaded = self.weights_loaded and not unpatch_weights
self.real_model = None
+ return True
+
+ def model_use_more_vram(self, extra_memory):
+ return self.model.partially_load(self.device, extra_memory)
def __eq__(self, other):
return self.model is other.model
+def use_more_memory(extra_memory, loaded_models, device):
+ for m in loaded_models:
+ if m.device == device:
+ extra_memory -= m.model_use_more_vram(extra_memory)
+ if extra_memory <= 0:
+ break
+
+def offloaded_memory(loaded_models, device):
+ offloaded_mem = 0
+ for m in loaded_models:
+ if m.device == device:
+ offloaded_mem += m.model_offloaded_memory()
+ return offloaded_mem
+
def minimum_inference_memory():
return (1024 * 1024 * 1024) * 1.2
@@ -365,6 +408,7 @@ def unload_model_clones(model, unload_weights_only=True, force_unload=True):
def free_memory(memory_required, device, keep_loaded=[]):
unloaded_model = []
can_unload = []
+ unloaded_models = []
for i in range(len(current_loaded_models) -1, -1, -1):
shift_model = current_loaded_models[i]
@@ -375,14 +419,18 @@ def free_memory(memory_required, device, keep_loaded=[]):
for x in sorted(can_unload):
i = x[-1]
+ memory_to_free = None
if not DISABLE_SMART_MEMORY:
- if get_free_memory(device) > memory_required:
+ free_mem = get_free_memory(device)
+ if free_mem > memory_required:
break
- current_loaded_models[i].model_unload()
- unloaded_model.append(i)
+ memory_to_free = memory_required - free_mem
+ logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
+ if current_loaded_models[i].model_unload(memory_to_free):
+ unloaded_model.append(i)
for i in sorted(unloaded_model, reverse=True):
- current_loaded_models.pop(i)
+ unloaded_models.append(current_loaded_models.pop(i))
if len(unloaded_model) > 0:
soft_empty_cache()
@@ -391,16 +439,17 @@ def free_memory(memory_required, device, keep_loaded=[]):
mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)
if mem_free_torch > mem_free_total * 0.25:
soft_empty_cache()
+ return unloaded_models
def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None):
global vram_state
inference_memory = minimum_inference_memory()
- extra_mem = max(inference_memory, memory_required)
+ extra_mem = max(inference_memory, memory_required) + 100 * 1024 * 1024
if minimum_memory_required is None:
minimum_memory_required = extra_mem
else:
- minimum_memory_required = max(inference_memory, minimum_memory_required)
+ minimum_memory_required = max(inference_memory, minimum_memory_required) + 100 * 1024 * 1024
models = set(models)
@@ -433,8 +482,16 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
devs = set(map(lambda a: a.device, models_already_loaded))
for d in devs:
if d != torch.device("cpu"):
- free_memory(extra_mem, d, models_already_loaded)
- return
+ free_memory(extra_mem + offloaded_memory(models_already_loaded, d), d, models_already_loaded)
+ free_mem = get_free_memory(d)
+ if free_mem < minimum_memory_required:
+ logging.info("Unloading models for lowram load.") #TODO: partial model unloading when this case happens, also handle the opposite case where models can be unlowvramed.
+ models_to_load = free_memory(minimum_memory_required, d)
+ logging.info("{} models unloaded.".format(len(models_to_load)))
+ else:
+ use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d)
+ if len(models_to_load) == 0:
+ return
logging.info(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
@@ -472,6 +529,14 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
cur_loaded_model = loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
current_loaded_models.insert(0, loaded_model)
+
+
+ devs = set(map(lambda a: a.device, models_already_loaded))
+ for d in devs:
+ if d != torch.device("cpu"):
+ free_mem = get_free_memory(d)
+ if free_mem > minimum_memory_required:
+ use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d)
return
@@ -541,7 +606,7 @@ def unet_inital_load_device(parameters, dtype):
return cpu_dev
def maximum_vram_for_weights(device=None):
- return (get_total_memory(device) * 0.8 - minimum_inference_memory())
+ return (get_total_memory(device) * 0.88 - minimum_inference_memory())
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
if args.bf16_unet:
@@ -567,12 +632,22 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor
if model_params * 2 > free_model_memory:
return fp8_dtype
- if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
- if torch.float16 in supported_dtypes:
- return torch.float16
- if should_use_bf16(device, model_params=model_params, manual_cast=True):
- if torch.bfloat16 in supported_dtypes:
- return torch.bfloat16
+ for dt in supported_dtypes:
+ if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params):
+ if torch.float16 in supported_dtypes:
+ return torch.float16
+ if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params):
+ if torch.bfloat16 in supported_dtypes:
+ return torch.bfloat16
+
+ for dt in supported_dtypes:
+ if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True):
+ if torch.float16 in supported_dtypes:
+ return torch.float16
+ if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True):
+ if torch.bfloat16 in supported_dtypes:
+ return torch.bfloat16
+
return torch.float32
# None means no manual cast
@@ -588,13 +663,13 @@ def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.flo
if bf16_supported and weight_dtype == torch.bfloat16:
return None
- if fp16_supported and torch.float16 in supported_dtypes:
- return torch.float16
+ for dt in supported_dtypes:
+ if dt == torch.float16 and fp16_supported:
+ return torch.float16
+ if dt == torch.bfloat16 and bf16_supported:
+ return torch.bfloat16
- elif bf16_supported and torch.bfloat16 in supported_dtypes:
- return torch.bfloat16
- else:
- return torch.float32
+ return torch.float32
def text_encoder_offload_device():
if args.gpu_only:
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index efac251ca..6c67193eb 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -1,8 +1,27 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
import torch
import copy
import inspect
import logging
import uuid
+import collections
import comfy.utils
import comfy.model_management
@@ -63,10 +82,31 @@ def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_
model_options["disable_cfg1_optimization"] = True
return model_options
+def wipe_lowvram_weight(m):
+ if hasattr(m, "prev_comfy_cast_weights"):
+ m.comfy_cast_weights = m.prev_comfy_cast_weights
+ del m.prev_comfy_cast_weights
+ m.weight_function = None
+ m.bias_function = None
+
+class LowVramPatch:
+ def __init__(self, key, model_patcher):
+ self.key = key
+ self.model_patcher = model_patcher
+ def __call__(self, weight):
+ return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
+
+
class ModelPatcher:
- def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
+ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
self.size = size
self.model = model
+ if not hasattr(self.model, 'device'):
+ logging.info("Model doesn't have a device attribute.")
+ self.model.device = offload_device
+ elif self.model.device is None:
+ self.model.device = offload_device
+
self.patches = {}
self.backup = {}
self.object_patches = {}
@@ -75,24 +115,32 @@ class ModelPatcher:
self.model_size()
self.load_device = load_device
self.offload_device = offload_device
- if current_device is None:
- self.current_device = self.offload_device
- else:
- self.current_device = current_device
-
self.weight_inplace_update = weight_inplace_update
- self.model_lowvram = False
- self.lowvram_patch_counter = 0
self.patches_uuid = uuid.uuid4()
+ if not hasattr(self.model, 'model_loaded_weight_memory'):
+ self.model.model_loaded_weight_memory = 0
+
+ if not hasattr(self.model, 'lowvram_patch_counter'):
+ self.model.lowvram_patch_counter = 0
+
+ if not hasattr(self.model, 'model_lowvram'):
+ self.model.model_lowvram = False
+
def model_size(self):
if self.size > 0:
return self.size
self.size = comfy.model_management.module_size(self.model)
return self.size
+ def loaded_size(self):
+ return self.model.model_loaded_weight_memory
+
+ def lowvram_patch_counter(self):
+ return self.model.lowvram_patch_counter
+
def clone(self):
- n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)
+ n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
n.patches = {}
for k in self.patches:
n.patches[k] = self.patches[k][:]
@@ -264,16 +312,16 @@ class ModelPatcher:
sd.pop(k)
return sd
- def patch_weight_to_device(self, key, device_to=None):
+ def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
if key not in self.patches:
return
weight = comfy.utils.get_attr(self.model, key)
- inplace_update = self.weight_inplace_update
+ inplace_update = self.weight_inplace_update or inplace_update
if key not in self.backup:
- self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)
+ self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
if device_to is not None:
temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
@@ -302,29 +350,24 @@ class ModelPatcher:
if device_to is not None:
self.model.to(device_to)
- self.current_device = device_to
+ self.model.device = device_to
+ self.model.model_loaded_weight_memory = self.model_size()
return self.model
- def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
- self.patch_model(device_to, patch_weights=False)
-
- logging.info("loading in lowvram mode {}".format(lowvram_model_memory/(1024 * 1024)))
- class LowVramPatch:
- def __init__(self, key, model_patcher):
- self.key = key
- self.model_patcher = model_patcher
- def __call__(self, weight):
- return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
-
+ def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
mem_counter = 0
patch_counter = 0
+ lowvram_counter = 0
for n, m in self.model.named_modules():
lowvram_weight = False
if hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m)
if mem_counter + module_mem >= lowvram_model_memory:
lowvram_weight = True
+ lowvram_counter += 1
+ if m.comfy_cast_weights:
+ continue
weight_key = "{}.weight".format(n)
bias_key = "{}.bias".format(n)
@@ -346,15 +389,37 @@ class ModelPatcher:
m.prev_comfy_cast_weights = m.comfy_cast_weights
m.comfy_cast_weights = True
else:
+ if hasattr(m, "comfy_cast_weights"):
+ if m.comfy_cast_weights:
+ wipe_lowvram_weight(m)
+
if hasattr(m, "weight"):
- self.patch_weight_to_device(weight_key, device_to)
- self.patch_weight_to_device(bias_key, device_to)
- m.to(device_to)
mem_counter += comfy.model_management.module_size(m)
+ param = list(m.parameters())
+ if len(param) > 0:
+ weight = param[0]
+ if weight.device == device_to:
+ continue
+
+ self.patch_weight_to_device(weight_key) #TODO: speed this up without OOM
+ self.patch_weight_to_device(bias_key)
+ m.to(device_to)
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
- self.model_lowvram = True
- self.lowvram_patch_counter = patch_counter
+ if lowvram_counter > 0:
+ logging.info("loaded in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024)))
+ self.model.model_lowvram = True
+ else:
+ logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024)))
+ self.model.model_lowvram = False
+ self.model.lowvram_patch_counter += patch_counter
+ self.model.device = device_to
+ self.model.model_loaded_weight_memory = mem_counter
+
+
+ def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
+ self.patch_model(device_to, patch_weights=False)
+ self.lowvram_load(device_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
return self.model
def calculate_weight(self, patches, weight, key):
@@ -527,34 +592,86 @@ class ModelPatcher:
def unpatch_model(self, device_to=None, unpatch_weights=True):
if unpatch_weights:
- if self.model_lowvram:
+ if self.model.model_lowvram:
for m in self.model.modules():
- if hasattr(m, "prev_comfy_cast_weights"):
- m.comfy_cast_weights = m.prev_comfy_cast_weights
- del m.prev_comfy_cast_weights
- m.weight_function = None
- m.bias_function = None
+ wipe_lowvram_weight(m)
- self.model_lowvram = False
- self.lowvram_patch_counter = 0
+ self.model.model_lowvram = False
+ self.model.lowvram_patch_counter = 0
keys = list(self.backup.keys())
- if self.weight_inplace_update:
- for k in keys:
- comfy.utils.copy_to_param(self.model, k, self.backup[k])
- else:
- for k in keys:
- comfy.utils.set_attr_param(self.model, k, self.backup[k])
+ for k in keys:
+ bk = self.backup[k]
+ if bk.inplace_update:
+ comfy.utils.copy_to_param(self.model, k, bk.weight)
+ else:
+ comfy.utils.set_attr_param(self.model, k, bk.weight)
self.backup.clear()
if device_to is not None:
self.model.to(device_to)
- self.current_device = device_to
+ self.model.device = device_to
+ self.model.model_loaded_weight_memory = 0
keys = list(self.object_patches_backup.keys())
for k in keys:
comfy.utils.set_attr(self.model, k, self.object_patches_backup[k])
self.object_patches_backup.clear()
+
+ def partially_unload(self, device_to, memory_to_free=0):
+ memory_freed = 0
+ patch_counter = 0
+
+ for n, m in list(self.model.named_modules())[::-1]:
+ if memory_to_free < memory_freed:
+ break
+
+ shift_lowvram = False
+ if hasattr(m, "comfy_cast_weights"):
+ module_mem = comfy.model_management.module_size(m)
+ weight_key = "{}.weight".format(n)
+ bias_key = "{}.bias".format(n)
+
+
+ if m.weight is not None and m.weight.device != device_to:
+ for key in [weight_key, bias_key]:
+ bk = self.backup.get(key, None)
+ if bk is not None:
+ if bk.inplace_update:
+ comfy.utils.copy_to_param(self.model, key, bk.weight)
+ else:
+ comfy.utils.set_attr_param(self.model, key, bk.weight)
+ self.backup.pop(key)
+
+ m.to(device_to)
+ if weight_key in self.patches:
+ m.weight_function = LowVramPatch(weight_key, self)
+ patch_counter += 1
+ if bias_key in self.patches:
+ m.bias_function = LowVramPatch(bias_key, self)
+ patch_counter += 1
+
+ m.prev_comfy_cast_weights = m.comfy_cast_weights
+ m.comfy_cast_weights = True
+ memory_freed += module_mem
+ logging.debug("freed {}".format(n))
+
+ self.model.model_lowvram = True
+ self.model.lowvram_patch_counter += patch_counter
+ self.model.model_loaded_weight_memory -= memory_freed
+ return memory_freed
+
+ def partially_load(self, device_to, extra_memory=0):
+ if self.model.model_lowvram == False:
+ return 0
+ if self.model.model_loaded_weight_memory + extra_memory > self.model_size():
+ pass #TODO: Full load
+ current_used = self.model.model_loaded_weight_memory
+ self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory)
+ return self.model.model_loaded_weight_memory - current_used
+
+ def current_loaded_device(self):
+ return self.model.device
diff --git a/comfy/sd.py b/comfy/sd.py
index fac1a487f..94fc4e590 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -564,7 +564,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
logging.debug("left over keys: {}".format(left_over))
if output_model:
- model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
+ model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
if inital_load_device != torch.device("cpu"):
logging.info("loaded straight to GPU")
model_management.load_model_gpu(model_patcher)
diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
index d32121d1b..e65cab285 100644
--- a/comfy/sd1_clip.py
+++ b/comfy/sd1_clip.py
@@ -313,6 +313,17 @@ def expand_directory_list(directories):
dirs.add(root)
return list(dirs)
+def bundled_embed(embed, prefix, suffix): #bundled embedding in lora format
+ i = 0
+ out_list = []
+ for k in embed:
+ if k.startswith(prefix) and k.endswith(suffix):
+ out_list.append(embed[k])
+ if len(out_list) == 0:
+ return None
+
+ return torch.cat(out_list, dim=0)
+
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
if isinstance(embedding_directory, str):
embedding_directory = [embedding_directory]
@@ -379,8 +390,12 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
elif embed_key is not None and embed_key in embed:
embed_out = embed[embed_key]
else:
- values = embed.values()
- embed_out = next(iter(values))
+ embed_out = bundled_embed(embed, 'bundle_emb.', '.string_to_param.*')
+ if embed_out is None:
+ embed_out = bundled_embed(embed, 'bundle_emb.', '.{}'.format(embed_key))
+ if embed_out is None:
+ values = embed.values()
+ embed_out = next(iter(values))
return embed_out
class SDTokenizer:
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 94fdcc0d2..d07a7106c 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -640,9 +640,9 @@ class Flux(supported_models_base.BASE):
unet_extra_config = {}
latent_format = latent_formats.Flux
- memory_usage_factor = 2.6
+ memory_usage_factor = 2.8
- supported_inference_dtypes = [torch.bfloat16, torch.float32]
+ supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
diff --git a/comfy/utils.py b/comfy/utils.py
index ec7d36607..e6736dbde 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -1,3 +1,22 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
+
import torch
import math
import struct
@@ -432,6 +451,12 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
+ k = "{}.attn.".format(prefix_from)
+ qkv = "{}.txt_attn.qkv.{}".format(prefix_to, end)
+ key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, hidden_size))
+ key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
+ key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
+
block_map = {"attn.to_out.0.weight": "img_attn.proj.weight",
"attn.to_out.0.bias": "img_attn.proj.bias",
}
diff --git a/comfy_extras/nodes_hunyuan.py b/comfy_extras/nodes_hunyuan.py
index a3ac8cb06..b03eaf6a2 100644
--- a/comfy_extras/nodes_hunyuan.py
+++ b/comfy_extras/nodes_hunyuan.py
@@ -19,6 +19,7 @@ class CLIPTextEncodeHunyuanDiT:
cond = output.pop("cond")
return ([[cond, output]], )
+
NODE_CLASS_MAPPINGS = {
"CLIPTextEncodeHunyuanDiT": CLIPTextEncodeHunyuanDiT,
}
diff --git a/comfy_extras/nodes_model_merging.py b/comfy_extras/nodes_model_merging.py
index b0d149c60..136f9a984 100644
--- a/comfy_extras/nodes_model_merging.py
+++ b/comfy_extras/nodes_model_merging.py
@@ -264,6 +264,7 @@ class CLIPSave:
metadata = {}
if not args.disable_metadata:
+ metadata["format"] = "pt"
metadata["prompt"] = prompt_info
if extra_pnginfo is not None:
for x in extra_pnginfo:
diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py
index ae9b85981..046096cba 100644
--- a/comfy_extras/nodes_sd3.py
+++ b/comfy_extras/nodes_sd3.py
@@ -100,3 +100,8 @@ NODE_CLASS_MAPPINGS = {
"CLIPTextEncodeSD3": CLIPTextEncodeSD3,
"ControlNetApplySD3": ControlNetApplySD3,
}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ # Sampling
+ "ControlNetApplySD3": "ControlNetApply SD3 and HunyuanDiT",
+}
diff --git a/execution.py b/execution.py
index 0a2e62e7e..d207e1b9e 100644
--- a/execution.py
+++ b/execution.py
@@ -188,6 +188,11 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
"current_inputs": input_data_formatted,
"current_outputs": output_data_formatted
}
+
+ if isinstance(ex, comfy.model_management.OOM_EXCEPTION):
+ logging.error("Got an OOM, unloading all loaded models.")
+ comfy.model_management.unload_all_models()
+
return (False, error_details, ex)
executed.add(unique_id)
diff --git a/folder_paths.py b/folder_paths.py
index 71faa2df4..3db1da61a 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -1,13 +1,13 @@
+from __future__ import annotations
+
import os
import time
import logging
-from typing import Set, List, Dict, Tuple
+from collections.abc import Collection
-supported_pt_extensions: Set[str] = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'])
+supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}
-SupportedFileExtensionsType = Set[str]
-ScanPathType = List[str]
-folder_names_and_paths: Dict[str, Tuple[ScanPathType, SupportedFileExtensionsType]] = {}
+folder_names_and_paths: dict[str, tuple[list[str], set[str]]] = {}
base_path = os.path.dirname(os.path.realpath(__file__))
models_dir = os.path.join(base_path, "models")
@@ -42,7 +42,7 @@ temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp
input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user")
-filename_list_cache = {}
+filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
if not os.path.exists(input_directory):
try:
@@ -50,33 +50,33 @@ if not os.path.exists(input_directory):
except:
logging.error("Failed to create input directory")
-def set_output_directory(output_dir):
+def set_output_directory(output_dir: str) -> None:
global output_directory
output_directory = output_dir
-def set_temp_directory(temp_dir):
+def set_temp_directory(temp_dir: str) -> None:
global temp_directory
temp_directory = temp_dir
-def set_input_directory(input_dir):
+def set_input_directory(input_dir: str) -> None:
global input_directory
input_directory = input_dir
-def get_output_directory():
+def get_output_directory() -> str:
global output_directory
return output_directory
-def get_temp_directory():
+def get_temp_directory() -> str:
global temp_directory
return temp_directory
-def get_input_directory():
+def get_input_directory() -> str:
global input_directory
return input_directory
#NOTE: used in http server so don't put folders that should not be accessed remotely
-def get_directory_by_type(type_name):
+def get_directory_by_type(type_name: str) -> str | None:
if type_name == "output":
return get_output_directory()
if type_name == "temp":
@@ -88,7 +88,7 @@ def get_directory_by_type(type_name):
# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format
# otherwise use default_path as base_dir
-def annotated_filepath(name):
+def annotated_filepath(name: str) -> tuple[str, str | None]:
if name.endswith("[output]"):
base_dir = get_output_directory()
name = name[:-9]
@@ -104,7 +104,7 @@ def annotated_filepath(name):
return name, base_dir
-def get_annotated_filepath(name, default_dir=None):
+def get_annotated_filepath(name: str, default_dir: str | None=None) -> str:
name, base_dir = annotated_filepath(name)
if base_dir is None:
@@ -116,7 +116,7 @@ def get_annotated_filepath(name, default_dir=None):
return os.path.join(base_dir, name)
-def exists_annotated_filepath(name):
+def exists_annotated_filepath(name) -> bool:
name, base_dir = annotated_filepath(name)
if base_dir is None:
@@ -126,17 +126,17 @@ def exists_annotated_filepath(name):
return os.path.exists(filepath)
-def add_model_folder_path(folder_name, full_folder_path):
+def add_model_folder_path(folder_name: str, full_folder_path: str) -> None:
global folder_names_and_paths
if folder_name in folder_names_and_paths:
folder_names_and_paths[folder_name][0].append(full_folder_path)
else:
folder_names_and_paths[folder_name] = ([full_folder_path], set())
-def get_folder_paths(folder_name):
+def get_folder_paths(folder_name: str) -> list[str]:
return folder_names_and_paths[folder_name][0][:]
-def recursive_search(directory, excluded_dir_names=None):
+def recursive_search(directory: str, excluded_dir_names: list[str] | None=None) -> tuple[list[str], dict[str, float]]:
if not os.path.isdir(directory):
return [], {}
@@ -153,6 +153,10 @@ def recursive_search(directory, excluded_dir_names=None):
logging.warning(f"Warning: Unable to access {directory}. Skipping this path.")
logging.debug("recursive file list on directory {}".format(directory))
+ dirpath: str
+ subdirs: list[str]
+ filenames: list[str]
+
for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True):
subdirs[:] = [d for d in subdirs if d not in excluded_dir_names]
for file_name in filenames:
@@ -160,7 +164,7 @@ def recursive_search(directory, excluded_dir_names=None):
result.append(relative_path)
for d in subdirs:
- path = os.path.join(dirpath, d)
+ path: str = os.path.join(dirpath, d)
try:
dirs[path] = os.path.getmtime(path)
except FileNotFoundError:
@@ -169,12 +173,12 @@ def recursive_search(directory, excluded_dir_names=None):
logging.debug("found {} files".format(len(result)))
return result, dirs
-def filter_files_extensions(files, extensions):
+def filter_files_extensions(files: Collection[str], extensions: Collection[str]) -> list[str]:
return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files)))
-def get_full_path(folder_name, filename):
+def get_full_path(folder_name: str, filename: str) -> str | None:
global folder_names_and_paths
if folder_name not in folder_names_and_paths:
return None
@@ -189,7 +193,7 @@ def get_full_path(folder_name, filename):
return None
-def get_filename_list_(folder_name):
+def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float]:
global folder_names_and_paths
output_list = set()
folders = folder_names_and_paths[folder_name]
@@ -199,9 +203,9 @@ def get_filename_list_(folder_name):
output_list.update(filter_files_extensions(files, folders[1]))
output_folders = {**output_folders, **folders_all}
- return (sorted(list(output_list)), output_folders, time.perf_counter())
+ return sorted(list(output_list)), output_folders, time.perf_counter()
-def cached_filename_list_(folder_name):
+def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float] | None:
global filename_list_cache
global folder_names_and_paths
if folder_name not in filename_list_cache:
@@ -222,7 +226,7 @@ def cached_filename_list_(folder_name):
return out
-def get_filename_list(folder_name):
+def get_filename_list(folder_name: str) -> list[str]:
out = cached_filename_list_(folder_name)
if out is None:
out = get_filename_list_(folder_name)
@@ -230,17 +234,17 @@ def get_filename_list(folder_name):
filename_list_cache[folder_name] = out
return list(out[0])
-def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):
- def map_filename(filename):
+def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, image_height=0) -> tuple[str, str, int, str, str]:
+ def map_filename(filename: str) -> tuple[int, str]:
prefix_len = len(os.path.basename(filename_prefix))
prefix = filename[:prefix_len + 1]
try:
digits = int(filename[prefix_len + 1:].split('_')[0])
except:
digits = 0
- return (digits, prefix)
+ return digits, prefix
- def compute_vars(input, image_width, image_height):
+ def compute_vars(input: str, image_width: int, image_height: int) -> str:
input = input.replace("%width%", str(image_width))
input = input.replace("%height%", str(image_height))
return input