diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 69917361a..72d1b8e4f 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -240,19 +240,32 @@ class QwenImagePipeline( prefix: str = "", ): super().__init__() + self.od_config = od_config self.device = get_local_device() model = od_config.model - self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler") + + # Check if model is a local path + local_files_only = os.path.exists(model) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) logger.info("Loaded Qwen-Image scheduler successfully") with set_default_torch_dtype(torch.bfloat16): - self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(model, subfolder="text_encoder") + self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) logger.info("Loaded Qwen-Image text encoder successfully") - self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").to(self.device) + self.vae = AutoencoderKLQwenImage.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self.device) logger.info("Loaded Qwen-Image VAE successfully") self.transformer = QwenImageTransformer2DModel() logger.info("Initialized Qwen-Image transformer successfully.") - self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer") + self.tokenizer = Qwen2Tokenizer.from_pretrained( + model, subfolder="tokenizer", local_files_only=local_files_only + ) logger.info("Loaded Qwen-Image tokenizer successfully.") self.stage = None @@ -763,10 +776,14 @@ class QwenImagePipeline( for name, tensor in state_dict.items(): yield name, tensor - model_name = "Qwen/Qwen-Image" try: - # Use the arguments as verified by the user - model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + # Get model path from config or download from HF + model_name = self.od_config.model if hasattr(self, "od_config") else "Qwen/Qwen-Image" + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + transformer_path = os.path.join(model_path, "transformer") self.transformer.load_weights(weight_iterator(transformer_path))