[Refactor][Diffusion] Move single-model layers into their model directories (#5908)

Signed-off-by: congw729 <115451386+congw729@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -41,13 +41,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_qkv_parallel_linear import (
     MoTQKVParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_row_parallel_linear import (
     MoTRowParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import (
     get_best_mot_config,
 )
 

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Layer-level correctness & performance test for MoTRMSNorm.
 
@@ -20,7 +23,7 @@ import torch
 from vllm import ir
 from vllm.config import VllmConfig, set_current_vllm_config
 
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 

tests/diffusion/kernels/mot/test_mot_norm_npu.py

@@ -1,8 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 
 import vllm_omni.diffusion.layers.custom_op as custom_op
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 

tests/diffusion/layers/test_custom_op.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the CustomOp platform-dispatch base class."""
 
 import pytest
@@ -39,7 +39,7 @@ def test_forward_cuda_still_requires_an_implementation():
 
 def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
     """MoTRMSNorm relies on the inherited fallback (regression test)."""
-    from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+    from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
     assert "forward_xpu" not in MoTRMSNorm.__dict__
 

tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for Wan2.2 quant_config propagation through transformer creation.
 
 Tests cover:
@@ -196,7 +196,7 @@ class TestPatchWanRmsNorm:
     """Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
 
     def test_patches_modules_with_wan_rms_norm(self):
-        from vllm_omni.diffusion.layers.norm import RMSNormVAE
+        from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
         from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
 
         # Create a fake module that has WanRMS_norm