hurukawa
3e705fcdb1
Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com> Signed-off-by: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com>
84 lines
3.4 KiB
Python
84 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import logging
|
|
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_classifier_free_guidance_world_size,
|
|
)
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class JoyImageEditCFGParallelMixin(CFGParallelMixin, ProgressBarMixin):
|
|
def cfg_normalize_function(self, noise_pred: torch.Tensor, comb_pred: torch.Tensor) -> torch.Tensor:
|
|
cond_norm = torch.norm(noise_pred, dim=2, keepdim=True)
|
|
comb_norm = torch.norm(comb_pred, dim=2, keepdim=True)
|
|
return comb_pred * (cond_norm / comb_norm.clamp_min(1e-6))
|
|
|
|
def diffuse(
|
|
self,
|
|
latents: torch.Tensor,
|
|
image_latents: torch.Tensor,
|
|
prompt_embeds: torch.Tensor,
|
|
prompt_embeds_mask: torch.Tensor,
|
|
negative_prompt_embeds: torch.Tensor | None,
|
|
negative_prompt_embeds_mask: torch.Tensor | None,
|
|
timesteps: torch.Tensor,
|
|
do_true_cfg: bool,
|
|
true_cfg_scale: float,
|
|
cfg_normalize: bool = True,
|
|
) -> torch.Tensor:
|
|
self.scheduler.set_begin_index(0)
|
|
with self.progress_bar(total=len(timesteps)) as pbar:
|
|
for timestep in timesteps:
|
|
if self.interrupt:
|
|
continue
|
|
self._current_timestep = timestep
|
|
|
|
latents = torch.cat([image_latents, latents[:, image_latents.shape[1] :]], dim=1)
|
|
latent_model_input = latents
|
|
timestep_expand = timestep.expand(latents.shape[0]).to(device=latents.device)
|
|
|
|
positive_kwargs = {
|
|
"hidden_states": latent_model_input,
|
|
"timestep": timestep_expand,
|
|
"encoder_hidden_states": prompt_embeds,
|
|
"encoder_hidden_states_mask": prompt_embeds_mask,
|
|
"return_dict": False,
|
|
}
|
|
negative_kwargs = None
|
|
if do_true_cfg:
|
|
negative_kwargs = {
|
|
"hidden_states": latent_model_input,
|
|
"timestep": timestep_expand,
|
|
"encoder_hidden_states": negative_prompt_embeds,
|
|
"encoder_hidden_states_mask": negative_prompt_embeds_mask,
|
|
"return_dict": False,
|
|
}
|
|
|
|
noise_pred = self.predict_noise_maybe_with_cfg(
|
|
do_true_cfg=do_true_cfg,
|
|
true_cfg_scale=true_cfg_scale,
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
cfg_normalize=cfg_normalize,
|
|
output_slice=None,
|
|
)
|
|
latents = self.scheduler_step_maybe_with_cfg(noise_pred, timestep, latents, do_true_cfg)
|
|
latents = torch.cat([image_latents, latents[:, image_latents.shape[1] :]], dim=1)
|
|
pbar.update()
|
|
return latents
|
|
|
|
def check_cfg_parallel_validity(self, true_cfg_scale: float) -> bool:
|
|
if get_classifier_free_guidance_world_size() == 1:
|
|
return True
|
|
if true_cfg_scale <= 1:
|
|
logger.warning("CFG parallel is enabled but Joy true_cfg_scale <= 1, so only the positive branch is used.")
|
|
return False
|
|
return True
|