From f764c6910d0d5857827cb07884da3f61cbc9b293 Mon Sep 17 00:00:00 2001 From: Dongjie Zou <85092850+baonudesifeizhai@users.noreply.github.com> Date: Tue, 2 Dec 2025 21:27:31 -0500 Subject: [PATCH] [diffusion] feat: support distilled vae generic (#14195) Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Mick --- .../configs/models/vaes/flux.py | 18 ++- .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/base.py | 25 ++++- .../configs/pipeline_configs/flux.py | 57 +++++++--- .../pipeline_configs/flux_finetuned.py | 103 ++++++++++++++++++ .../runtime/loader/component_loader.py | 66 ++++++----- .../pipelines_core/composed_pipeline_base.py | 15 ++- .../runtime/pipelines_core/stages/decoding.py | 28 ++++- .../multimodal_gen/runtime/server_args.py | 18 ++- .../runtime/utils/hf_diffusers_utils.py | 47 ++++---- 10 files changed, 302 insertions(+), 79 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py diff --git a/python/sglang/multimodal_gen/configs/models/vaes/flux.py b/python/sglang/multimodal_gen/configs/models/vaes/flux.py index 53bd44618..33308640a 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/flux.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/flux.py @@ -47,9 +47,21 @@ class FluxVAEConfig(VAEConfig): ) * 2 def post_init(self): - self.arch_config.vae_scale_factor = 2 ** ( - len(self.arch_config.block_out_channels) - 1 - ) + # Calculate vae_scale_factor: prefer block_out_channels, fallback to dim_mult or scale_factor_spatial + if ( + hasattr(self.arch_config, "block_out_channels") + and self.arch_config.block_out_channels + ): + self.arch_config.vae_scale_factor = 2 ** ( + len(self.arch_config.block_out_channels) - 1 + ) + elif self.arch_config.dim_mult: + self.arch_config.vae_scale_factor = 2 ** ( + len(self.arch_config.dim_mult) - 1 + ) + else: + self.arch_config.vae_scale_factor = self.arch_config.scale_factor_spatial + self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 250c0d7ab..c482cc4a7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -5,6 +5,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( SlidingTileAttnConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( + Flux2FinetunedPipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( FastHunyuanConfig, HunyuanConfig, @@ -23,6 +26,7 @@ __all__ = [ "HunyuanConfig", "FastHunyuanConfig", "FluxPipelineConfig", + "Flux2FinetunedPipelineConfig", "PipelineConfig", "SlidingTileAttnConfig", "WanT2V480PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 5ffc92783..5d3437247 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -281,7 +281,7 @@ class PipelineConfig: return image_latents # called after scale_and_shift, before vae decoding - def preprocess_decoding(self, latents): + def preprocess_decoding(self, latents, server_args=None, vae=None): return latents def gather_latents_for_sp(self, latents): @@ -487,9 +487,30 @@ class PipelineConfig: raise ValueError("model_path is required in kwargs") # 1. Get the pipeline config class from the registry + from sglang.multimodal_gen.configs.pipeline_configs.flux import ( + Flux2PipelineConfig, + ) + model_info = get_model_info(model_path) - pipeline_config = model_info.pipeline_config_cls() + # 1.5. Adjust pipeline config for fine-tuned VAE if needed + pipeline_config_cls = model_info.pipeline_config_cls + vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path") + + # Check if this is a Flux2 model with fal/FLUX.2-Tiny-AutoEncoder + if ( + isinstance(pipeline_config_cls, type) + and issubclass(pipeline_config_cls, Flux2PipelineConfig) + and vae_path is not None + and "FLUX.2-Tiny-AutoEncoder" in vae_path + ): + from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( + Flux2FinetunedPipelineConfig, + ) + + pipeline_config_cls = Flux2FinetunedPipelineConfig + + pipeline_config = pipeline_config_cls() # 2. Load PipelineConfig from a json file or a PipelineConfig object if provided if isinstance(pipeline_config_or_path, str): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index d4bed105c..99a058e9a 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -215,10 +215,10 @@ def _prepare_latent_ids( t = torch.arange(1) # [0] - time dimension h = torch.arange(height) w = torch.arange(width) - l = torch.arange(1) # [0] - layer dimension + layer = torch.arange(1) # [0] - layer dimension # Create position IDs: (H*W, 4) - latent_ids = torch.cartesian_prod(t, h, w, l) + latent_ids = torch.cartesian_prod(t, h, w, layer) # Expand to batch: (B, H*W, 4) latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1) @@ -289,9 +289,9 @@ def _prepare_text_ids( t = torch.arange(1) if t_coord is None else t_coord[i] h = torch.arange(1) w = torch.arange(1) - l = torch.arange(L) + layer = torch.arange(L) - coords = torch.cartesian_prod(t, h, w, l) + coords = torch.cartesian_prod(t, h, w, layer) out_ids.append(coords) return torch.stack(out_ids) @@ -500,7 +500,6 @@ class Flux2PipelineConfig(FluxPipelineConfig): return image_latents def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): - txt_ids = _prepare_text_ids(prompt_embeds).to(device=device) img_ids = batch.latent_ids @@ -550,19 +549,51 @@ class Flux2PipelineConfig(FluxPipelineConfig): image_latents = _patchify_latents(image_latents) return image_latents - def preprocess_decoding(self, latents): - latents = _unpatchify_latents(latents) + def _check_vae_has_bn(self, vae): + """Check if VAE has bn attribute (cached check to avoid repeated hasattr calls).""" + if not hasattr(self, "_vae_has_bn_cache"): + self._vae_has_bn_cache = hasattr(vae, "bn") and vae.bn is not None + return self._vae_has_bn_cache + + def preprocess_decoding(self, latents, server_args=None, vae=None): + """Preprocess latents before decoding. + + Dynamically adapts based on VAE type: + - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels) + - Distilled VAE (no bn): keeps patchified latents (128 channels) + """ + if vae is not None and self._check_vae_has_bn(vae): + return _unpatchify_latents(latents) return latents def get_decode_scale_and_shift(self, device, dtype, vae): + """Get scale and shift for decoding. + + Dynamically adapts based on VAE type: + - Standard Flux2 VAE (has bn): uses BatchNorm statistics + - Distilled VAE (no bn): uses scaling_factor from config + """ vae_arch_config = self.vae_config.arch_config - latents_bn_mean = ( - vae.bn.running_mean.view(1, -1, 1, 1).to(device=device).to(device, dtype) + + if self._check_vae_has_bn(vae): + # Standard Flux2 VAE: use BatchNorm statistics + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype) + latents_bn_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps + ).to(device, dtype) + return 1 / latents_bn_std, latents_bn_mean + + # Distilled VAE or unknown: use scaling_factor + scaling_factor = ( + getattr(vae.config, "scaling_factor", None) + if hasattr(vae, "config") + else getattr(vae, "scaling_factor", None) + ) or getattr(vae_arch_config, "scaling_factor", 0.13025) + + scale = torch.tensor(scaling_factor, device=device, dtype=dtype).view( + 1, 1, 1, 1 ) - latents_bn_std = torch.sqrt( - vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps - ).to(device, dtype) - return 1 / latents_bn_std, latents_bn_mean + return 1 / scale, None def post_denoising_loop(self, latents, batch): latent_ids = batch.latent_ids diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py new file mode 100644 index 000000000..ebafdb5be --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py @@ -0,0 +1,103 @@ +""" +Pipeline configuration for Flux fine-tuned/distilled models. + +This module provides specialized handling for Flux fine-tuned models from HuggingFace, +such as fal/FLUX.2-Tiny-AutoEncoder and other community fine-tuned variants. + +Key differences from standard Flux2PipelineConfig: +- Handles custom VAE architectures loaded via auto_map +- Supports both patchified (128 channels) and unpatchified (32 channels) latents +- Dynamically adapts scale/shift based on VAE type +- Properly handles 5D latents (batch, channels, frames, height, width) for decoding +""" + +from dataclasses import dataclass + +import torch + +from sglang.multimodal_gen.configs.pipeline_configs.flux import ( + Flux2PipelineConfig, + _unpatchify_latents, +) + + +@dataclass +class Flux2FinetunedPipelineConfig(Flux2PipelineConfig): + """ + Pipeline configuration for Flux fine-tuned/distilled models. + + This configuration automatically detects and handles custom VAE architectures + (e.g., Flux2TinyAutoEncoder) loaded via HuggingFace's auto_map mechanism. + + Features: + - Automatic VAE type detection (standard vs. distilled) + - Proper handling of patchified/unpatchified latents + - Support for custom scaling factors from fine-tuned models + - 5D latents support for both single-frame and multi-frame generation + """ + + def preprocess_decoding( + self, latents: torch.Tensor, server_args=None, vae=None + ) -> torch.Tensor: + """ + Preprocess latents before decoding. + + Handles both standard Flux2 VAE and fine-tuned/distilled VAEs: + - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels) + - Distilled/Finetuned VAE (no bn): keeps patchified latents (128 channels) + + Also handles 5D latents (batch, channels, frames, height, width) by converting + to 4D (batch, channels, height, width) for single-frame cases. + + Args: + latents: Input latents tensor, can be 4D or 5D + server_args: Server arguments (optional, for compatibility) + vae: VAE model instance for dynamic type detection + + Returns: + Preprocessed latents ready for VAE decoding + """ + # Handle 5D latents (batch, channels, frames, height, width) + if latents.ndim == 5: + batch_size, channels, frames, height, width = latents.shape + if frames == 1: + latents = latents.squeeze(2) + else: + latents = latents.permute(0, 2, 1, 3, 4).contiguous() + latents = latents.view(batch_size * frames, channels, height, width) + + if vae is not None and self._check_vae_has_bn(vae): + latents = _unpatchify_latents(latents) + return latents + + def get_decode_scale_and_shift(self, device, dtype, vae): + """ + Get scale and shift for decoding. + + Dynamically adapts based on VAE type: + - Standard Flux2 VAE (has bn): uses BatchNorm statistics + - Distilled/Finetuned VAE (no bn): uses scaling_factor from config + + Args: + device: Target device for tensors + dtype: Target dtype for tensors + vae: VAE model instance + + Returns: + Tuple of (scaling_factor, shift_factor) + - scaling_factor: Tensor or scalar to divide latents by + - shift_factor: Tensor or scalar to add to latents (None for distilled VAEs) + """ + vae_arch_config = self.vae_config.arch_config + + if self._check_vae_has_bn(vae): + # Standard Flux2 VAE: use BatchNorm statistics + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype) + latents_bn_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps + ).to(device, dtype) + return 1 / latents_bn_std, latents_bn_mean + + # Distilled/Finetuned VAE: Flux2TinyAutoEncoder doesn't need external scaling + scale = torch.tensor(1.0, device=device, dtype=dtype).view(1, 1, 1, 1) + return scale, None diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 3e93e80a5..9ae7abf11 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -4,10 +4,11 @@ import dataclasses import glob +import importlib.util import json import os import time -from abc import ABC, abstractmethod +from abc import ABC from collections.abc import Generator, Iterable from copy import deepcopy from typing import Any, cast @@ -101,17 +102,13 @@ def load_native(library, component_module_path: str, server_args: ServerArgs): revision=server_args.revision, ) elif library == "diffusers": - import diffusers + from diffusers import AutoModel - config = get_diffusers_component_config(model_path=component_module_path) - class_name = config.pop("_class_name", None) - if class_name: - cls = getattr(diffusers, class_name) - return cls.from_pretrained( - component_module_path, revision=server_args.revision, **config - ) - else: - raise ValueError("Cannot determine class name for generic diffusers loader") + return AutoModel.from_pretrained( + component_module_path, + revision=server_args.revision, + trust_remote_code=server_args.trust_remote_code, + ) else: raise ValueError(f"Unsupported library: {library}") @@ -202,13 +199,6 @@ class ComponentLoader(ABC): f"load_customized not implemented for {self.__class__.__name__}" ) - @abstractmethod - def load_customized( - self, model_path: str, server_args: ServerArgs, module_name: str - ) -> Any: - """Implement the minimal core load logic in subclasses.""" - raise NotImplementedError - @classmethod def for_module_type( cls, module_type: str, transformers_or_diffusers: str @@ -503,7 +493,6 @@ class ImageEncoderLoader(TextEncoderLoader): encoder_config.update_model_arch(model_config) # Always start with local device; load_model will adjust for offload if needed - should_offload = self.should_offload(server_args) # TODO(will): add support for other dtypes return self.load_model( component_model_path, @@ -555,7 +544,7 @@ class VAELoader(ComponentLoader): ): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_component_config(model_path=component_model_path) - class_name = config.pop("_class_name") + class_name = config.pop("_class_name", None) assert ( class_name is not None ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." @@ -571,23 +560,42 @@ class VAELoader(ComponentLoader): target_device = self.target_device(server_args.vae_cpu_offload) - with set_default_torch_dtype( - PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - ), skip_init_modules(): + # Check for auto_map first (custom VAE classes) + auto_map = config.get("auto_map", {}) + auto_model_map = auto_map.get("AutoModel") + if auto_model_map: + module_path, cls_name = auto_model_map.rsplit(".", 1) + custom_module_file = os.path.join(component_model_path, f"{module_path}.py") + spec = importlib.util.spec_from_file_location("_custom", custom_module_file) + custom_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(custom_module) + vae_cls = getattr(custom_module, cls_name) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + with set_default_torch_dtype(vae_dtype): + vae = vae_cls.from_pretrained( + component_model_path, + revision=server_args.revision, + trust_remote_code=server_args.trust_remote_code, + ) + vae = vae.to(device=target_device, dtype=vae_dtype) + return vae.eval() + + # Load from ModelRegistry (standard VAE classes) + with ( + set_default_torch_dtype( + PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + ), + skip_init_modules(), + ): vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) vae = vae_cls(vae_config).to(target_device) - # Find all safetensors files safetensors_list = _list_safetensors_files(component_model_path) - # TODO(PY) assert ( len(safetensors_list) == 1 ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}" loaded = safetensors_load_file(safetensors_list[0]) - vae.load_state_dict( - loaded, strict=False - ) # We might only load encoder or decoder - + vae.load_state_dict(loaded, strict=False) return vae.eval() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 175c14fe8..d6cb99c1b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -281,7 +281,6 @@ class ComposedPipelineBase(ABC): transformers_or_diffusers, architecture, ) in tqdm(iterable=model_index.items(), desc="Loading required modules"): - if transformers_or_diffusers is None: logger.warning( "Module %s in model_index.json has null value, removing from required_config_modules", @@ -304,7 +303,19 @@ class ComposedPipelineBase(ABC): else: load_module_name = module_name - component_model_path = os.path.join(self.model_path, load_module_name) + # Use custom VAE path if provided, otherwise use default path + if module_name == "vae" and server_args.vae_path is not None: + component_model_path = server_args.vae_path + # Download from HuggingFace Hub if path doesn't exist locally + if not os.path.exists(component_model_path): + component_model_path = maybe_download_model(component_model_path) + logger.info( + "Using custom VAE path: %s instead of default path: %s", + component_model_path, + os.path.join(self.model_path, load_module_name), + ) + else: + component_model_path = os.path.join(self.model_path, load_module_name) module = PipelineComponentLoader.load_module( module_name=load_module_name, component_model_path=component_model_path, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 2d9cf7f07..bf5932b14 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -27,6 +27,26 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) +def _ensure_tensor_decode_output(decode_output): + """ + Ensure VAE decode output is a tensor. + + Some VAE implementations return DecoderOutput objects with a .sample attribute, + tuples, or tensors directly. This function normalizes the output to always be a tensor. + + Args: + decode_output: Output from VAE.decode(), can be DecoderOutput, tuple, or torch.Tensor + + Returns: + torch.Tensor: The decoded image tensor + """ + if isinstance(decode_output, tuple): + return decode_output[0] + if hasattr(decode_output, "sample"): + return decode_output.sample + return decode_output + + class DecodingStage(PipelineStage): """ Stage for decoding latent representations into pixel space. @@ -106,7 +126,10 @@ class DecodingStage(PipelineStage): # scale and shift latents = self.scale_and_shift(latents, server_args) - latents = server_args.pipeline_config.preprocess_decoding(latents) + # Preprocess latents before decoding (e.g., unpatchify for standard Flux2 VAE) + latents = server_args.pipeline_config.preprocess_decoding( + latents, server_args, vae=self.vae + ) # Decode latents with torch.autocast( @@ -120,7 +143,8 @@ class DecodingStage(PipelineStage): pass if not vae_autocast_enabled: latents = latents.to(vae_dtype) - image = self.vae.decode(latents) + decode_output = self.vae.decode(latents) + image = _ensure_tensor_decode_output(decode_output) # De-normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 4abc5cbfc..c95199b50 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py """The arguments of sglang-diffusion Inference.""" + import argparse import dataclasses import inspect @@ -252,6 +253,9 @@ class ServerArgs: # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. lora_path: str | None = None lora_nickname: str = "default" # for swapping adapters in the pipeline + + # VAE parameters + vae_path: str | None = None # Custom VAE path (e.g., for distilled autoencoder) # can restrict layers to adapt, e.g. ["q_proj"] # Will adapt only q, k, v, o by default. lora_target_modules: list[str] | None = None @@ -374,6 +378,12 @@ class ServerArgs: type=str, help="Directory containing StepVideo model", ) + parser.add_argument( + "--vae-path", + type=str, + default=ServerArgs.vae_path, + help="Custom path to VAE model (e.g., for distilled autoencoder). If not specified, VAE will be loaded from the main model path.", + ) # attention parser.add_argument( @@ -825,17 +835,15 @@ class ServerArgs: if self.ulysses_degree is None: self.ulysses_degree = 1 logger.info( - f"Ulysses degree not set, " f"using default value {self.ulysses_degree}" + f"Ulysses degree not set, using default value {self.ulysses_degree}" ) if self.ring_degree is None: self.ring_degree = 1 - logger.info( - f"Ring degree not set, " f"using default value {self.ring_degree}" - ) + logger.info(f"Ring degree not set, using default value {self.ring_degree}") if self.ring_degree > 1: - if self.attention_backend != None and self.attention_backend != "fa": + if self.attention_backend is not None and self.attention_backend != "fa": raise ValueError( "Ring Attention is only supported for flash attention backend for now" ) diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index a3c83b6f7..e0d04f272 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -135,33 +135,33 @@ def get_diffusers_component_config( """Gets a configuration of a submodule for the given diffusers model. Args: - model_path: the path of the submodule + model_path: the path of the submodule (can be local path or HuggingFace model ID) Returns: The loaded configuration. """ - # Check if the model path exists - if os.path.exists(model_path): - # tokenizer - config_names = ["generation_config.json"] - # By default, we load config.json, but scheduler_config.json for scheduler - if "scheduler" in model_path: - config_names.append("scheduler_config.json") - else: - config_names.append("config.json") + # Download from HuggingFace Hub if path doesn't exist locally + if not os.path.exists(model_path): + model_path = maybe_download_model(model_path) - config_file_paths = [ - os.path.join(model_path, config_name) for config_name in config_names - ] - - combined_config = reduce( - lambda acc, path: acc | load_dict(path), config_file_paths, {} - ) - - return combined_config + # tokenizer + config_names = ["generation_config.json"] + # By default, we load config.json, but scheduler_config.json for scheduler + if "scheduler" in model_path: + config_names.append("scheduler_config.json") else: - raise RuntimeError(f"Diffusers config file not found at {model_path}") + config_names.append("config.json") + + config_file_paths = [ + os.path.join(model_path, config_name) for config_name in config_names + ] + + combined_config = reduce( + lambda acc, path: acc | load_dict(path), config_file_paths, {} + ) + + return combined_config # Models don't use the same configuration key for determining the maximum @@ -390,9 +390,10 @@ def maybe_download_model( logger.info( "Downloading model snapshot from HF Hub for %s...", model_name_or_path ) - with get_lock(model_name_or_path).acquire( - poll_interval=2 - ), suppress_other_loggers(not_suppress_on_main_rank=True): + with ( + get_lock(model_name_or_path).acquire(poll_interval=2), + suppress_other_loggers(not_suppress_on_main_rank=True), + ): local_path = snapshot_download( repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"],