[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)

python/sglang/multimodal_gen/test/run_suite.py

@@ -90,7 +90,7 @@ def collect_test_items(files, filter_expr=None):
         cmd.extend(["-k", filter_expr])
     cmd.extend(files)
 
-    logger.info(f"Collecting tests with command: {' '.join(cmd)}")
+    print(f"Collecting tests with command: {' '.join(cmd)}")
     result = subprocess.run(cmd, capture_output=True, text=True)
 
     # Check for collection errors
@@ -114,7 +114,7 @@ def collect_test_items(files, filter_expr=None):
         raise RuntimeError(error_msg)
 
     if result.returncode == 5:
-        logger.info(
+        print(
             "No tests were collected (exit code 5). This may be expected with filters."
         )
 
@@ -130,7 +130,7 @@ def collect_test_items(files, filter_expr=None):
             if "::" in test_id:
                 test_items.append(test_id)
 
-    logger.info(f"Collected {len(test_items)} test items")
+    print(f"Collected {len(test_items)} test items")
     return test_items
 
 
@@ -145,20 +145,21 @@ def run_pytest(files, filter_expr=None):
     if filter_expr:
         base_cmd.extend(["-k", filter_expr])
 
-    max_retries = 4
+    max_retries = 6
     # retry if the perf assertion failed, for {max_retries} times
     for i in range(max_retries + 1):
         cmd = list(base_cmd)
         if i > 0:
             cmd.append("--last-failed")
-        cmd.extend(files)
+        else:
+            cmd.extend(files)
 
         if i > 0:
-            logger.info(
+            print(
                 f"Performance assertion failed. Retrying ({i}/{max_retries}) with --last-failed..."
             )
 
-        logger.info(f"Running command: {' '.join(cmd)}")
+        print(f"Running command: {' '.join(cmd)}")
 
         process = subprocess.Popen(
             cmd,
@@ -185,7 +186,7 @@ def run_pytest(files, filter_expr=None):
         # Exit code 5 means no tests were collected/selected - treat as success
         # when using filters, since some partitions may have all tests filtered out
         if returncode == 5:
-            logger.info(
+            print(
                 "No tests collected (exit code 5). This is expected when filters "
                 "deselect all tests in a partition. Treating as success."
             )
@@ -210,7 +211,7 @@ def run_pytest(files, filter_expr=None):
         if not (is_perf_assertion or is_flaky_ci_assertion or is_oom_error):
             return returncode
 
-    logger.info(f"Max retry exceeded")
+    print(f"Max retry exceeded")
     return returncode
 
 
@@ -261,12 +262,13 @@ def main():
     print(f"Selected {len(suite_files_abs)} files:")
     for f in suite_files_abs:
         print(f"  - {os.path.basename(f)}")
-    print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}")
 
     if not my_items:
         print("No items assigned to this partition. Exiting success.")
         sys.exit(0)
 
+    print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}")
+
     # 4. execute with the specific test items
     exit_code = run_pytest(my_items)
     sys.exit(exit_code)