diff --git a/extensions_built_in/diffusion_models/flux3/flux3.py b/extensions_built_in/diffusion_models/flux3/flux3.py index 3be9205..ae09c38 100644 --- a/extensions_built_in/diffusion_models/flux3/flux3.py +++ b/extensions_built_in/diffusion_models/flux3/flux3.py @@ -10,11 +10,16 @@ t2i and i2i (reference images as clean conditioning tokens). in an ``AdvancedPromptEmbeds`` (padded to a multiple of 80 tokens -- the DiT has no text mask, pad tokens are part of the trained context). -``model.name_or_path`` is the model package (hub repo id or local dir): +``model.name_or_path`` is the model package (hub repo id, default +``black-forest-labs/flux-3-image-klein-base``, or a local dir): - config.json model.safetensors ae.safetensors text_encoder/ tokenizer/ + config.json model.safetensors ae.safetensors text_encoder/ -Only the DiT and the image VAE come from the package. The package's +Only the DiT and the image VAE come from the package, and they live in the +ComfyUI layout under MODELS_PATH: ``diffusion_models/flux_3_image_klein_base.safetensors`` +and ``vae/flux_3_vae.safetensors``. A file already there (or anywhere under its +category folder) is used in place; only a missing one is downloaded from the +package repo and renamed to that path. The package's text_encoder/ and tokenizer/ are byte-identical to Qwen/Qwen3-VL-4B-Instruct, so the text encoder loads from the Comfy-Org single file under MODELS_PATH/text_encoders/ (downloaded there once when missing) with config @@ -41,6 +46,7 @@ from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel +from toolkit.paths import MODELS_PATH from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) @@ -66,6 +72,13 @@ if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO +DEFAULT_REPO = "black-forest-labs/flux-3-image-klein-base" +# package file -> ComfyUI-layout location under MODELS_PATH +COMFY_FILES = { + "model.safetensors": "diffusion_models/flux_3_image_klein_base.safetensors", + "ae.safetensors": "vae/flux_3_vae.safetensors", +} + # Static flow-matching timeshift used for the training timestep/noise # schedule (the package's sampler alpha). DEFAULT_TIMESTEP_SHIFT = DEFAULT_ALPHA @@ -105,6 +118,7 @@ class Flux3Model(BaseModel): self.model_config.model_kwargs.get("text_pad_multiple", DEFAULT_PAD_MULTIPLE) ) self.package_config = None + self._package_config_checked = False # references are area-matched to the target by default; off = keep their # own size, only scaling down when larger than the target area self.match_target_res = bool( @@ -126,7 +140,8 @@ class Flux3Model(BaseModel): # ------------------------------------------------------------------ def _package_file(self, filename: str, required: bool = True) -> Optional[str]: """A package file: model_kwargs ``_path`` override, else the file - inside a local package dir, else downloaded from the hub repo.""" + inside a local package dir, else its ComfyUI-layout spot under + MODELS_PATH (weights), else downloaded from the hub repo.""" override = self.model_config.model_kwargs.get( filename.split(".")[0] + "_path", None ) @@ -144,6 +159,8 @@ class Flux3Model(BaseModel): return None if os.path.isfile(base) and filename == "model.safetensors": return base # name_or_path points straight at the DiT file + if filename in COMFY_FILES: + return self._comfy_weight(filename, base) from huggingface_hub import hf_hub_download try: @@ -156,12 +173,46 @@ class Flux3Model(BaseModel): self.print_and_status_update(f" - no {filename} in {base} ({type(e).__name__})") return None + def _comfy_weight(self, filename: str, repo_id: str) -> str: + """``MODELS_PATH//`` (or the local name anywhere + under the category folder); a miss downloads the package file and + renames it to that spot.""" + from toolkit.models.v2.resolver import find_file_recursive + + rel_path = COMFY_FILES[filename] + target = os.path.join(MODELS_PATH, rel_path) + if os.path.exists(target): + return target + found = find_file_recursive( + os.path.join(MODELS_PATH, os.path.dirname(rel_path)), + os.path.basename(rel_path), + ) + if found is not None: + return found + from huggingface_hub import hf_hub_download + + self.print_and_status_update( + f"Downloading {filename} from {repo_id} to {target}" + ) + os.makedirs(os.path.dirname(target), exist_ok=True) + path = hf_hub_download( + repo_id=repo_id, + filename=filename, + token=os.getenv("HF_TOKEN", None), + local_dir=os.path.dirname(target), + ) + os.replace(path, target) + return target + def _load_package_config(self) -> Optional[dict]: - if self.package_config is None: + if not self._package_config_checked: + self._package_config_checked = True path = self._package_file("config.json", required=False) if path is not None: with open(path, "r") as f: self.package_config = json.load(f) + else: + self.print_and_status_update(" - using the canonical image-klein DiT layout") return self.package_config def _transformer_params(self) -> Flux3TransformerParams: diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 28afd31..b818927 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1717,13 +1717,13 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ }, { name: "flux3", - label: "FLUX.3 (small)", + label: "FLUX.3 Image Klein (base)", group: "image", - gateUrl: "https://huggingface.co/herzogenhorn/f3-image-small", + gateUrl: "https://huggingface.co/black-forest-labs/flux-3-image-klein-base", defaults: { // default updates when [selected, unselected] in the UI "config.process[0].model.name_or_path": [ - "herzogenhorn/f3-image-small", + "black-forest-labs/flux-3-image-klein-base", defaultNameOrPath, ], "config.process[0].model.quantize": [true, false], diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 1b4edc2..4a89167 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -55,7 +55,7 @@ ARCH_REGISTRY: Dict[str, dict] = { "wan22_14b_i2v": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}, "needs_control_image": True}, "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, - "flux3": {"modality": "image", "model": {"name_or_path": "herzogenhorn/f3-image-small", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}}, + "flux3": {"modality": "image", "model": {"name_or_path": "black-forest-labs/flux-3-image-klein-base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}}, "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}}, "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}}, "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},