[diffusion] improve: skip loading vision module for text encoders (#16304)

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -23,6 +23,9 @@ from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig
+from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+    QwenImageEditPipelineConfig,
+)
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.fsdp_load import (
     maybe_load_fsdp_model,
@@ -462,6 +465,14 @@ class TextEncoderLoader(ComponentLoader):
             with local_torch_device, skip_init_modules():
                 architectures = getattr(model_config, "architectures", [])
                 model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+                enable_image_understanding = (
+                    True
+                    if isinstance(
+                        server_args.pipeline_config, QwenImageEditPipelineConfig
+                    )
+                    else False
+                )
+                model_config.enable_image_understanding = enable_image_understanding
                 model = model_cls(model_config)
 
             weights_to_load = {name for name, _ in model.named_parameters()}

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -1,7 +1,5 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
-from types import SimpleNamespace
-
 from transformers import (
     Cache,
     DynamicCache,
@@ -508,13 +506,15 @@ class Qwen2_5_VLModel(nn.Module):
     accepts_loss_kwargs = False
     _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
 
-    def __init__(self, config):
+    def __init__(self, config, enable_image_understanding: bool = False):
         super().__init__()
-        self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(
-            config.vision_config
-        )
         self.language_model = Qwen2_5_VLTextModel(config.text_config)
-        self.visual.to(torch.get_default_dtype())
+
+        if enable_image_understanding:
+            self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(
+                config.vision_config
+            )
+            self.visual.to(torch.get_default_dtype())
         self.rope_deltas = None  # cache rope_deltas here
         self.config = config
         # Initialize weights and apply final processing
@@ -997,41 +997,6 @@ class Qwen2_5_VLModel(nn.Module):
         return output if return_dict else output.to_tuple()
 
 
-class DotDict(dict):
-    def __init__(self, mapping):
-        super().__init__()
-        for key, value in mapping.items():
-            if isinstance(value, dict):
-                value = DotDict(value)  # 递归转换
-            elif isinstance(value, list):
-                # 如果是 list，且元素是 dict 也递归转换
-                value = [
-                    DotDict(item) if isinstance(item, dict) else item for item in value
-                ]
-            self[key] = value
-
-    def __getattr__(self, item):
-        try:
-            return self[item]
-        except KeyError:
-            raise AttributeError(f"No attribute '{item}'")
-
-    def __setattr__(self, key, value):
-        self[key] = value
-
-    def __delattr__(self, key):
-        del self[key]
-
-
-def dict_to_namespace(d):
-    for k, v in d.items():
-        if isinstance(v, dict):
-            d[k] = dict_to_namespace(v)
-        elif isinstance(v, list):
-            d[k] = [dict_to_namespace(i) if isinstance(i, dict) else i for i in v]
-    return SimpleNamespace(**d)
-
-
 class Qwen2_5_VLForConditionalGeneration(TextEncoder):
     # BitandBytes specific attributes
     default_bitsandbytes_target_modules = [
@@ -1058,12 +1023,17 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder):
         prefix: str = "",
     ) -> None:
         super().__init__(config)
+        enable_image_understanding = config.enable_image_understanding
         config = config.arch_config
-        self.model = Qwen2_5_VLModel(config)
+        self.model = Qwen2_5_VLModel(
+            config, enable_image_understanding=enable_image_understanding
+        )
         self.lm_head = nn.Linear(
             config.text_config.hidden_size, config.text_config.vocab_size, bias=False
         )
 
+        self.enable_image_understanding = enable_image_understanding
+
         self.config = config
 
     def get_input_embeddings(self):
@@ -1157,6 +1127,8 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder):
 
             name = name.replace("model.", "model.language_model.")
             if "visual." in name:
+                if not self.enable_image_understanding:
+                    continue
                 name = name.replace("visual.", "model.visual.")
             try:
                 # Skip loading extra bias for GPTQ models.
@@ -1164,7 +1136,6 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder):
                     continue
                 param = params_dict[name]
             except KeyError:
-                print(params_dict.keys())
                 raise
 
             weight_loader = getattr(param, "weight_loader", default_weight_loader)