diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index ea17e6858..9e477063d 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -36,7 +36,17 @@ def parse_args() -> argparse.Namespace: "--model", default="Qwen/Qwen-Image", help="Diffusion model name or local path. Supported models: " - "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1", + "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, " + "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, " + "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, " + "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, " + "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to a YAML file containing stage configurations for Omni.", ) parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.") parser.add_argument( @@ -311,6 +321,8 @@ def main(): **lora_args, **quant_kwargs, } + if args.stage_configs_path: + omni_kwargs["stage_configs_path"] = args.stage_configs_path if use_nextstep: # NextStep-1.1 requires explicit pipeline class omni_kwargs["model_class_name"] = "NextStep11Pipeline" @@ -338,6 +350,8 @@ def main(): print(f" Image size: {args.width}x{args.height}") if args.lora_path: print(f" LoRA: scale={args.lora_scale}") + if args.stage_configs_path: + print(f" stage-configs-path: {args.stage_configs_path}") print(f"{'=' * 60}\n") # Build LoRA request when --lora-path is set diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index daccef343..e2c6e2341 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -328,8 +328,15 @@ class DiffusersPipelineLoader: # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights + # NOTE: if the model is quantized, ignore not_loaded check for scale weights + weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")} + weights_not_loaded = weights_not_loaded - weights_scale_not_loaded if weights_not_loaded: raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") + if weights_scale_not_loaded: + logger.warning( + f"Following weight_scale weights were not initialized from checkpoint: {weights_scale_not_loaded}" + ) def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool: quant_config = od_config.quantization_config diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index a89931550..7987844ad 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -1565,7 +1565,9 @@ class HunYuanAttention(nn.Module): class HunyuanImage3DecoderLayer(nn.Module): - def __init__(self, config: HunyuanImage3Config, layer_idx: int, prefix: str = ""): + def __init__( + self, config: HunyuanImage3Config, quant_config: QuantizationConfig | None, layer_idx: int, prefix: str = "" + ): super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx @@ -1607,9 +1609,13 @@ class HunyuanImage3DecoderLayer(nn.Module): (isinstance(config.num_experts, int) and config.num_experts > 1) or (isinstance(config.num_experts, list) and max(config.num_experts) > 1) ) and layer_idx >= config.moe_layer_num_skipped: - self.mlp = HunYuanSparseMoeBlock(config, layer_id=layer_idx, prefix=f"{prefix}.mlp") + self.mlp = HunYuanSparseMoeBlock( + config, quant_config=quant_config, layer_id=layer_idx, prefix=f"{prefix}.mlp" + ) else: - self.mlp = HunYuanMLP(self.hidden_size, self.intermediate_size, config.hidden_act) + self.mlp = HunYuanMLP( + self.hidden_size, self.intermediate_size, config.hidden_act, quant_config=quant_config + ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1704,14 +1710,15 @@ class HunyuanImage3PreTrainedModel(PreTrainedModel): class HunyuanImage3Model(nn.Module): - def __init__(self, config: HunyuanImage3Config, prefix: str = ""): + def __init__(self, config: HunyuanImage3Config, quant_config: QuantizationConfig | None, prefix: str = ""): super().__init__() - quant_config = None lora_config = None self.num_redundant_experts = 0 self.config = config self.device = get_local_device() + self.quant_config = quant_config + logger.debug(f"quant_config: {quant_config}") self.padding_idx = config.pad_token_id lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0 self.vocab_size = config.vocab_size + lora_vocab @@ -1730,6 +1737,7 @@ class HunyuanImage3Model(nn.Module): config.num_hidden_layers, lambda prefix: HunyuanImage3DecoderLayer( config=config, + quant_config=quant_config, layer_idx=int(prefix.split(".")[-1]), prefix=prefix, ), diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py index e4b717a69..2fe486714 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py @@ -20,6 +20,7 @@ from vllm.transformers_utils.config import get_config from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from .autoencoder import AutoencoderKLConv3D @@ -77,7 +78,8 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): fall_back_to_pt=True, ) ] - self.model = HunyuanImage3Model(self.hf_config) + quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config) + self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config) self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae) self._pipeline = None self._tkwrapper = TokenizerWrapper(od_config.model)