[Feature] Diffusion LoRA Adapter Support (PEFT compatible) for vLLM alignment (#758)

Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk>
Signed-off-by: Andy Zhou <jzhoubc@connect.ust.hk>
Signed-off-by: dongbo910220 <1275604947@qq.com>
Signed-off-by: Andy Zhou <46011930+AndyZhou952@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: dongbo910220 <1275604947@qq.com>
Co-authored-by: Samit <285365963@qq.com>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -94,6 +94,7 @@ class DiffusersPipelineLoader:
         load_format = self.load_config.load_format
         use_safetensors = False
         index_file = DIFFUSION_MODEL_WEIGHTS_INDEX
+        index_file_with_subfolder = f"{subfolder}/{index_file}" if subfolder else index_file
 
         # only hf is supported currently
         if load_format == "auto":
@@ -129,8 +130,8 @@ class DiffusersPipelineLoader:
         for pattern in allow_patterns:
             hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
             if len(hf_weights_files) > 0:
-                if pattern == "*.safetensors":
-                    use_safetensors = True
+                # Decide by actual files rather than pattern name (patterns may include subfolders).
+                use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files)
                 break
 
         if use_safetensors:
@@ -142,11 +143,22 @@ class DiffusersPipelineLoader:
             if not is_local:
                 download_safetensors_index_file_from_hf(
                     model_name_or_path,
-                    index_file,
+                    index_file_with_subfolder,
                     self.load_config.download_dir,
                     revision,
                 )
-            hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file)
+            # Some diffusers pipelines keep component weights under a
+            # subfolder (e.g. "transformer/") and the corresponding index file
+            # uses filenames relative to that subfolder. vLLM's
+            # `filter_duplicate_safetensors_files` expects weight_map entries
+            # to be relative to the `hf_folder` we pass in, so we point it to
+            # the component subfolder to avoid filtering out all shards.
+            filter_folder = os.path.join(hf_folder, subfolder) if subfolder is not None else hf_folder
+            hf_weights_files = filter_duplicate_safetensors_files(
+                hf_weights_files,
+                filter_folder,
+                index_file,
+            )
         else:
             hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
 
@@ -188,8 +200,9 @@ class DiffusersPipelineLoader:
 
     def download_model(self, model_config: ModelConfig) -> None:
         self._prepare_weights(
-            model_config.model,
-            model_config.revision,
+            model_name_or_path=model_config.model,
+            subfolder=None,
+            revision=model_config.revision,
             fall_back_to_pt=True,
             allow_patterns_overrides=None,
         )