[BugFix]: fix a lot of bug (#1565)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

Author: princepride

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -128,13 +128,10 @@ def __init__(self, vision_model):
         else:
             self.vision_model = vision_model
 
-        # Configure weights for linear equivalent of patch embedding
-        self.patch_embed_weight = self.vision_model.embeddings.patch_embedding.weight
-        self.patch_embed_bias = self.vision_model.embeddings.patch_embedding.bias
-
     def forward(self, packed_pixel_values, packed_flattened_position_ids, cu_seqlens, max_seqlen):
-        w = self.patch_embed_weight.view(self.patch_embed_weight.shape[0], -1)
-        x = F.linear(packed_pixel_values, w, self.patch_embed_bias)
+        patch_embed = self.vision_model.embeddings.patch_embedding
+        w = patch_embed.weight.view(patch_embed.weight.shape[0], -1)
+        x = F.linear(packed_pixel_values, w, patch_embed.bias)
         pos = self.vision_model.embeddings.position_embedding(packed_flattened_position_ids)
         x = x + pos
         hidden_states = x.unsqueeze(0)
@@ -205,6 +202,9 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         )
         vit_config_path = os.path.join(model_path, "vit_config.json")
         vit_conf = SiglipVisionConfig.from_json_file(vit_config_path)
+        if vit_conf.num_hidden_layers == 27:
+            vit_conf.num_hidden_layers = 26
+        vit_conf.vision_use_head = False
         self.vit_model = SiglipVisionModel(vit_conf)
         self.image_processor = SiglipImageProcessor.from_pretrained(model_path, local_files_only=True)
 

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -871,8 +871,8 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load transformer weights."""
-        # Filter weights for transformer only
         transformer_weights = (
-            (name.replace("transformer.", ""), weight) for name, weight in weights if name.startswith("transformer.")
+            (name.replace("transformer.", "", 1), weight) for name, weight in weights if name.startswith("transformer.")
         )
-        return self.transformer.load_weights(transformer_weights)
+        loaded = self.transformer.load_weights(transformer_weights)
+        return {f"transformer.{name}" for name in loaded}

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -86,7 +86,7 @@ def _get_data_with_lock(self, lock_file: str, shm_handle: dict):
             return obj, int(shm_handle.get("size", 0))
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
-            return None, 0
+            return None
         finally:
             # If data has been received, delete lock_file.
             if obj and os.path.exists(lock_file):
@@ -105,32 +105,32 @@ def get(
                 metadata = metadata.get(get_key)
 
             if not isinstance(metadata, dict):
-                return None, 0
+                return None
 
             if "inline_bytes" in metadata:
                 try:
                     obj = self.deserialize_obj(metadata["inline_bytes"])
                     return obj, int(metadata.get("size", 0))
                 except Exception as e:
                     logger.error(f"SharedMemoryConnector inline get failed for req {get_key}: {e}")
-                    return None, 0
+                    return None
 
             if "shm" in metadata:
                 shm_handle = metadata["shm"]
                 lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock"
                 return self._get_data_with_lock(lock_file, shm_handle)
 
-            return None, 0
+            return None
         shm = None
         try:
             shm = shm_pkg.SharedMemory(name=get_key)
             if shm is None or shm.size == 0:
-                return None, 0
+                return None
             lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock"
             shm_handle = {"name": get_key, "size": shm.size}
             return self._get_data_with_lock(lock_file, shm_handle)
         except Exception:
-            return None, 0
+            return None
         finally:
             if shm:
                 shm.close()

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -136,6 +136,8 @@ def _poll_single_request(self, request: Request):
             logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}")
             return False
 
+        if result is None:
+            return False
         payload_data, size = result
 
         if payload_data:

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -36,12 +36,12 @@
     GlmImageVQVAEConfig,
 )
 from transformers.models.glm_image.processing_glm_image import GlmImageProcessor
-from vllm.attention.layer import Attention
 from vllm.config import CacheConfig, MultiModalConfig, VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.distributed import utils as dist_utils
 from vllm.logger import init_logger
+from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention.mm_encoder_attention import (
     MMEncoderAttention,
 )