[Refactor][Diffusion] Read the resolved offload policy outside the config layer (#7327)
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
22 files changed
+265
-111
No files matched your search
@@ -33,6 +33,10 @@ from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
|
||||
convert_diffusers_quantization_config,
|
||||
ensure_supported_diffusers_quantization,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.config import (
|
||||
OffloadStrategy,
|
||||
resolve_offload_strategy,
|
||||
)
|
||||
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt
|
||||
@@ -148,9 +152,10 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
self._accept_call_kwargs = set(inspect.signature(self._pipeline.__call__).parameters.keys())
|
||||
|
||||
# CPU offloading
|
||||
if self.od_config.enable_layerwise_offload:
|
||||
strategy = resolve_offload_strategy(self.od_config)
|
||||
if strategy is OffloadStrategy.LAYER_WISE:
|
||||
self._pipeline.enable_sequential_cpu_offload()
|
||||
elif self.od_config.enable_cpu_offload:
|
||||
elif strategy is OffloadStrategy.MODEL_LEVEL:
|
||||
self._pipeline.enable_model_cpu_offload()
|
||||
|
||||
# VAE slicing and tiling: try-catch because not all models have VAE
|
||||
|
||||
Reference in new issue
Block a user