[diffusion][MOVA] fix: fix task type in MOVA pipeline and shared model placement (#19489)
This commit is contained in:
committed by
GitHub
parent
5acb45cf32
commit
4ec450e97b
@@ -27,7 +27,7 @@ logger = init_logger(__name__)
|
||||
class MOVAPipelineConfig(PipelineConfig):
|
||||
"""Configuration for MOVA (text+image -> video+audio) pipelines."""
|
||||
|
||||
task_type: ModelTaskType = ModelTaskType.T2V
|
||||
task_type: ModelTaskType = ModelTaskType.I2V
|
||||
|
||||
# Model configs
|
||||
dit_config: MOVAVideoConfig = field(default_factory=MOVAVideoConfig)
|
||||
|
||||
@@ -19,6 +19,8 @@ class MOVASamplingParams(SamplingParams):
|
||||
visual_shift: float = 5.0
|
||||
audio_shift: float = 5.0
|
||||
|
||||
adjust_frames: bool = False
|
||||
|
||||
negative_prompt: str = (
|
||||
"色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,"
|
||||
"整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,"
|
||||
|
||||
@@ -220,7 +220,7 @@ class ConditionalCrossAttention(nn.Module):
|
||||
head_size=self.head_dim,
|
||||
causal=False,
|
||||
softmax_scale=None,
|
||||
is_cross_attention=True,
|
||||
# is_cross_attention=True,
|
||||
)
|
||||
|
||||
def forward(
|
||||
|
||||
@@ -349,6 +349,11 @@ class MOVADenoisingStage(PipelineStage):
|
||||
self._manage_device_placement(current_model, model_to_offload, server_args)
|
||||
return current_model
|
||||
|
||||
def _ensure_shared_models_on_device(self, server_args: ServerArgs):
|
||||
"""Ensure shared denoising modules are on the active device when cpu offload is enabled."""
|
||||
self._manage_device_placement(self.audio_dit, None, server_args)
|
||||
self._manage_device_placement(self.dual_tower_bridge, None, server_args)
|
||||
|
||||
def _apply_guidance_rescale(
|
||||
self,
|
||||
noise_pred,
|
||||
@@ -378,7 +383,7 @@ class MOVADenoisingStage(PipelineStage):
|
||||
@torch.no_grad()
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
self._maybe_compile_dits(server_args)
|
||||
self._manage_device_placement(self.audio_dit, None, server_args)
|
||||
self._ensure_shared_models_on_device(server_args)
|
||||
|
||||
paired_timesteps = batch.paired_timesteps
|
||||
if paired_timesteps is None:
|
||||
|
||||
Reference in New Issue
Block a user