github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Models v2 - phase 0

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 10:51:43 -06:00
parent 5497a001cb
commit e8d9cf6d35
29 changed files with 583 additions and 401 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
extensions_built_in/diffusion_models/boogu_image/boogu_image.py
View File

@@ -446,14 +446,5 @@ class BooguImageModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["double_stream_layers", "single_stream_layers"]
def convert_lora_weights_before_save(self, state_dict):
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}

15
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -443,21 +443,8 @@ class ChromaModel(BaseModel):
batch = kwargs.get('batch')
return (noise - batch.latents).detach()
def convert_lora_weights_before_save(self, state_dict):
# currently starte with transformer. but needs to start with diffusion_model. for comfyui
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
# saved as diffusion_model. but needs to be transformer. for ai-toolkit
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
def get_base_model_version(self):
return "chroma"

15
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -425,21 +425,8 @@ class ChromaRadianceModel(BaseModel):
batch = kwargs.get('batch')
return (noise - batch.latents).detach()
def convert_lora_weights_before_save(self, state_dict):
# currently starte with transformer. but needs to start with diffusion_model. for comfyui
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
# saved as diffusion_model. but needs to be transformer. for ai-toolkit
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
def get_base_model_version(self):
return "chroma_radiance"

13
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
View File

@@ -377,16 +377,5 @@ class ErnieImageModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd

19
extensions_built_in/diffusion_models/example_model/example_model.py
View File

@@ -504,19 +504,8 @@ class ExampleModel(BaseModel):
attribute in src/model.py."""
return ["blocks"]
def convert_lora_weights_before_save(self, state_dict):
"""Map internal LoRA keys to the ecosystem-standard naming right before
the .safetensors is written. Most modern models ship LoRAs with a
``diffusion_model.`` prefix (ComfyUI convention); internally ai-toolkit
uses ``transformer.``."""
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
# LoRA keys save with the ecosystem-standard ``diffusion_model.`` prefix
# (ComfyUI convention) and load back to the internal ``transformer.``
# prefix; see BaseModel.convert_lora_weights_before_save/load
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
"""Inverse of the above, applied when resuming from a saved LoRA."""
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}

15
extensions_built_in/diffusion_models/f_light/f_light.py
View File

@@ -270,21 +270,8 @@ class FLiteModel(BaseModel):
# return (noise - batch.latents).detach()
return (batch.latents - noise).detach()
def convert_lora_weights_before_save(self, state_dict):
# currently starte with transformer. but needs to start with diffusion_model. for comfyui
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
# saved as diffusion_model. but needs to be transformer. for ai-toolkit
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
def get_base_model_version(self):
return "f-lite"

14
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

@@ -505,19 +505,7 @@ class Flux2Model(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["double_blocks", "single_blocks"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:

15
extensions_built_in/diffusion_models/hidream/hidream_model.py
View File

@@ -432,21 +432,8 @@ class HidreamModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ['double_stream_blocks', 'single_stream_blocks']
def convert_lora_weights_before_save(self, state_dict):
# currently starte with transformer. but needs to start with diffusion_model. for comfyui
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
# saved as diffusion_model. but needs to be transformer. for ai-toolkit
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
def get_base_model_version(self):
return "hidream_i1"

13
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
View File

@@ -620,16 +620,5 @@ class Ideogram4Model(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd

11
extensions_built_in/diffusion_models/krea2/krea2.py
View File

@@ -867,14 +867,5 @@ class Krea2Model(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["blocks"]
def convert_lora_weights_before_save(self, state_dict):
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}

110
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -1191,23 +1191,17 @@ class LTX2Model(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["transformer_blocks"]
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
new_sd = convert_lora_diffusers_to_original(new_sd, version=self.ltx_version)
return new_sd
state_dict = super().convert_lora_weights_before_save(state_dict)
return convert_lora_diffusers_to_original(state_dict, version=self.ltx_version)
def convert_lora_weights_before_load(self, state_dict):
state_dict = convert_lora_original_to_diffusers(
state_dict, version=self.ltx_version
)
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
return super().convert_lora_weights_before_load(state_dict)
class LTX23Model(LTX2Model):
@@ -1239,88 +1233,30 @@ class LTX25Model(LTX2Model):
ltx_te_path = None
# ------------------------------------------------------------------
# ComfyUI-style file resolution (mirrors MinimaxH3Model)
# ComfyUI-style file resolution (toolkit/models/v2/resolver.py)
# ------------------------------------------------------------------
@staticmethod
def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]:
if not os.path.isdir(root_dir):
return None
for dirpath, dirnames, filenames in os.walk(root_dir):
dirnames.sort()
if filename in filenames:
return os.path.join(dirpath, filename)
return None
def _resolve_comfy_file(self, component: str) -> str:
"""Find a weight file at its local location, or download it there
when (and only when) it is missing.
Search order: model_kwargs override, the repo-relative path under
MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare
filename at the root, any subfolder of the component's category
folder (recursive), then the hub — downloaded to the repo-relative
path under MODELS_PATH.
"""
override = self.model_config.model_kwargs.get(f"{component}_path", None)
if override is not None:
if not os.path.exists(override):
raise FileNotFoundError(
f"model_kwargs.{component}_path does not exist: {override}"
)
return override
rel_path = COMFY_LTX25_FILES[component]
filename = os.path.basename(rel_path)
category = os.path.dirname(rel_path)
for rel in (rel_path, filename):
candidate = os.path.join(MODELS_PATH, rel)
if os.path.exists(candidate):
return candidate
found = self._find_file_recursive(os.path.join(MODELS_PATH, category), filename)
if found is not None:
return found
repo_id = COMFY_LTX25_REPO
name_or_path = self.model_config.name_or_path
if (
name_or_path
and not os.path.exists(name_or_path)
and not name_or_path.endswith(".safetensors")
and "/" in name_or_path
):
repo_id = name_or_path
self.print_and_status_update(
f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}"
from toolkit.models.v2.resolver import (
repo_id_from_name_or_path,
resolve_comfy_file,
)
return huggingface_hub.hf_hub_download(
repo_id=repo_id, filename=rel_path, token=HF_TOKEN, local_dir=MODELS_PATH
return resolve_comfy_file(
COMFY_LTX25_FILES[component],
repo_id=repo_id_from_name_or_path(
self.model_config.name_or_path, COMFY_LTX25_REPO
),
override_path=self.model_config.model_kwargs.get(
f"{component}_path", None
),
hf_token=HF_TOKEN,
status_fn=self.print_and_status_update,
)
def _resolve_named_file(self, path: str, component: str) -> str:
"""Resolve an explicit .safetensors path: local file, models-folder
file, or an 'org/repo/path/file.safetensors' hub path (downloaded
into the models folder)."""
if os.path.exists(path):
return path
splits = path.split("/")
if len(splits) < 3:
raise ValueError(
f"Invalid {component} path: {path}. Must be a local file or "
"'repo_id/repo/filename.safetensors' to download from the Hugging Face Hub."
)
rel_path = "/".join(splits[2:])
for candidate in (
os.path.join(MODELS_PATH, rel_path),
os.path.join(MODELS_PATH, splits[-1]),
):
if os.path.exists(candidate):
return candidate
return huggingface_hub.hf_hub_download(
repo_id="/".join(splits[:2]),
filename=rel_path,
token=HF_TOKEN,
local_dir=MODELS_PATH,
)
from toolkit.models.v2.resolver import resolve_named_file
return resolve_named_file(path, component=component, hf_token=HF_TOKEN)
def _resolve_dit_path(self) -> str:
name_or_path = self.model_config.name_or_path

13
extensions_built_in/diffusion_models/mageflow/mageflow.py
View File

@@ -629,18 +629,7 @@ class MageFlowModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["transformer_blocks"]
def convert_lora_weights_before_save(self, state_dict):
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}
lora_keys_use_comfy_prefix = True
class MageFlowEditModel(MageFlowModel):
arch = "mageflow_edit"

93
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -52,6 +52,11 @@ from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.memory_management import MemoryManager
from toolkit.metadata import get_meta_for_safetensors
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.resolver import (
find_file_recursive,
repo_id_from_name_or_path,
resolve_comfy_file,
)
from toolkit.paths import MODELS_PATH
from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
from toolkit.util.ostris_quant import OstrisLinear
@@ -233,64 +238,23 @@ class MinimaxH3Model(BaseModel):
# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------
@staticmethod
def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]:
"""First (breadth-stable, sorted) match of ``filename`` anywhere under
``root_dir``."""
if not os.path.isdir(root_dir):
return None
for dirpath, dirnames, filenames in os.walk(root_dir):
dirnames.sort()
if filename in filenames:
return os.path.join(dirpath, filename)
return None
def _resolve_comfy_file(self, component: str) -> str:
"""Find a weight file at its local location, or download it there
when (and only when) it is missing.
Search order: model_kwargs override, the repo-relative path under
MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare
filename at the root, any subfolder of the component's category
folder (recursive — e.g. diffusion_models/my_custom_sub/), the same
spots under name_or_path when it is a local folder, then the hub —
downloaded to the repo-relative path under MODELS_PATH.
"""
override = self.model_config.model_kwargs.get(f"{component}_path", None)
if override is not None:
if not os.path.exists(override):
raise FileNotFoundError(
f"model_kwargs.{component}_path does not exist: {override}"
)
return override
rel_path = COMFY_FILES[component]
filename = os.path.basename(rel_path)
category = os.path.dirname(rel_path)
roots = [MODELS_PATH]
"""Find a weight file at its local location (model_kwargs override,
comfy layout under MODELS_PATH or a local name_or_path dir), or
download it there when (and only when) it is missing — see
toolkit/models/v2/resolver.py for the search order."""
name_or_path = self.model_config.name_or_path
if name_or_path and os.path.isdir(name_or_path):
roots.append(name_or_path)
for root in roots:
for rel in (rel_path, filename):
candidate = os.path.join(root, rel)
if os.path.exists(candidate):
return candidate
for root in roots:
found = self._find_file_recursive(os.path.join(root, category), filename)
if found is not None:
return found
import huggingface_hub
repo_id = COMFY_REPO
if name_or_path and not os.path.exists(name_or_path) and "/" in name_or_path:
repo_id = name_or_path
self.print_and_status_update(
f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}"
extra_roots = (
[name_or_path] if name_or_path and os.path.isdir(name_or_path) else []
)
return huggingface_hub.hf_hub_download(
repo_id=repo_id, filename=rel_path, local_dir=MODELS_PATH
return resolve_comfy_file(
COMFY_FILES[component],
repo_id=repo_id_from_name_or_path(name_or_path, COMFY_REPO),
override_path=self.model_config.model_kwargs.get(
f"{component}_path", None
),
extra_roots=extra_roots,
status_fn=self.print_and_status_update,
)
def _dit_component(self) -> str:
@@ -322,7 +286,7 @@ class MinimaxH3Model(BaseModel):
lora_path = self.model_config.assistant_lora_path
if not os.path.exists(lora_path):
filename = os.path.basename(lora_path)
found = self._find_file_recursive(
found = find_file_recursive(
os.path.join(MODELS_PATH, "loras"), filename
)
if found is not None:
@@ -1226,20 +1190,9 @@ class MinimaxH3Model(BaseModel):
"*adaln_proj*",
]
def convert_lora_weights_before_save(self, state_dict):
# ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the
# standard diffusion_model prefix maps directly
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}
# ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the
# standard diffusion_model prefix maps directly
lora_keys_use_comfy_prefix = True
class MinimaxH3Ref2VAModel(MinimaxH3Model):
"""Reference-to-video (ref2va): the control images ride along as reference

13
extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py
View File

@@ -405,16 +405,5 @@ class NucleusImageModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["transformer_blocks"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd

16
extensions_built_in/diffusion_models/omnigen2/__init__.py
View File

@@ -343,21 +343,7 @@ class OmniGen2Model(BaseModel):
return ["noise_refiner", "context_refiner", "ref_image_refiner", "layers"]
return ["noise_refiner", "context_refiner", "layers"]
def convert_lora_weights_before_save(self, state_dict):
# currently starte with transformer. but needs to start with diffusion_model. for comfyui
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
def convert_lora_weights_before_load(self, state_dict):
# saved as diffusion_model. but needs to be transformer. for ai-toolkit
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def get_base_model_version(self):
return "omnigen2"

11
extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py
View File

@@ -332,14 +332,5 @@ class PRXPixelT2IModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["blocks"]
def convert_lora_weights_before_save(self, state_dict):
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}

14
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -415,19 +415,7 @@ class QwenImageModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["transformer_blocks"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:

17
extensions_built_in/diffusion_models/z_image/z_image.py
View File

@@ -31,8 +31,8 @@ try:
from diffusers import ZImagePipeline
# our subclass of the diffusers transformer with the universal loading /
# quantization mixin (see toolkit/models/classes/_mixin.py)
from toolkit.models.v2.z_image import ZImageTransformer2DModel
# quantization mixin (see toolkit/models/v2/_mixin.py)
from toolkit.models.v2.diffusion_models.z_image import ZImageTransformer2DModel
except ImportError:
raise ImportError(
"Diffusers is out of date. Update diffusers to the latest version by doing pip uninstall diffusers and then pip install -r requirements.txt"
@@ -464,16 +464,5 @@ class ZImageModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd

14
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -27,7 +27,6 @@ from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, va
from .zeta_chroma_pipeline import ZetaChromaPipeline
scheduler_config = {
"num_train_timesteps": 1000,
"use_dynamic_shifting": False,
@@ -368,16 +367,5 @@ class ZetaChromaModel(BaseModel):
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
lora_keys_use_comfy_prefix = True
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd