563 lines
22 KiB
Python
563 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py
|
|
|
|
import json
|
|
import os
|
|
from collections.abc import Callable, Iterable
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
|
|
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
|
|
FlowMatchEulerDiscreteScheduler,
|
|
)
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import ErnieImageTransformer2DModel
|
|
from vllm_omni.diffusion.models.interface import SupportImageInput
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
|
|
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _resolve_model_path_for_optional_pe(model: str, revision: str | None) -> str:
|
|
if os.path.exists(model):
|
|
return model
|
|
try:
|
|
return download_weights_from_hf_specific(model, revision, ["pe/*", "pe_tokenizer/*"])
|
|
except Exception as e:
|
|
logger.debug("Failed to resolve ERNIE-Image PE files for %s: %s", model, e)
|
|
return model
|
|
|
|
|
|
def get_ernie_image_post_process_func(od_config: OmniDiffusionConfig):
|
|
if od_config.output_type == "latent":
|
|
return lambda x: x
|
|
model_name = od_config.model
|
|
if os.path.exists(model_name):
|
|
model_path = model_name
|
|
else:
|
|
model_path = download_weights_from_hf_specific(model_name, None, ["*"])
|
|
|
|
vae_config_path = os.path.join(model_path, "vae/config.json")
|
|
with open(vae_config_path) as f:
|
|
vae_config = json.load(f)
|
|
vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
|
|
|
|
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_normalize=False)
|
|
|
|
def post_process_func(images: torch.Tensor):
|
|
images = (images.clamp(-1, 1) + 1) / 2
|
|
return image_processor.postprocess(images)
|
|
|
|
return post_process_func
|
|
|
|
|
|
class ErnieImagePipeline(
|
|
nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
|
|
):
|
|
support_image_input = False
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
od_config: OmniDiffusionConfig,
|
|
prefix: str = "",
|
|
is_distilled: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
self.is_distilled = is_distilled
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=True,
|
|
)
|
|
]
|
|
|
|
self._execution_device = get_local_device()
|
|
model = od_config.model
|
|
logger.info("Model path for initialization: %s", model)
|
|
local_files_only = os.path.exists(model)
|
|
logger.info("Local files only: %s", local_files_only)
|
|
|
|
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
model,
|
|
subfolder="scheduler",
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
self.text_encoder = AutoModel.from_pretrained(
|
|
model,
|
|
subfolder="text_encoder",
|
|
torch_dtype=od_config.dtype,
|
|
local_files_only=local_files_only,
|
|
).to(self._execution_device)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
model,
|
|
subfolder="tokenizer",
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
self.vae = AutoencoderKLFlux2.from_pretrained(
|
|
model,
|
|
subfolder="vae",
|
|
torch_dtype=od_config.dtype,
|
|
local_files_only=local_files_only,
|
|
).to(self._execution_device)
|
|
|
|
# Load PE (Prompt Enhancement) model if available. For repo IDs,
|
|
# resolve/download only PE files first so the existence check works.
|
|
pe_base_path = _resolve_model_path_for_optional_pe(model, getattr(od_config, "revision", None))
|
|
pe_model_path = os.path.join(pe_base_path, "pe")
|
|
if os.path.exists(pe_model_path):
|
|
try:
|
|
self.pe_model = AutoModelForCausalLM.from_pretrained(
|
|
pe_model_path,
|
|
torch_dtype=od_config.dtype,
|
|
local_files_only=True,
|
|
trust_remote_code=True,
|
|
).to(self._execution_device)
|
|
self.pe_tokenizer = AutoTokenizer.from_pretrained(
|
|
pe_base_path,
|
|
subfolder="pe_tokenizer",
|
|
local_files_only=True,
|
|
trust_remote_code=True,
|
|
use_fast=False,
|
|
)
|
|
self.use_pe = True
|
|
logger.info("Loaded PE model from %s", pe_model_path)
|
|
except Exception as e:
|
|
logger.warning("Failed to load PE model: %s", e)
|
|
self.pe_model = None
|
|
self.pe_tokenizer = None
|
|
self.use_pe = False
|
|
else:
|
|
self.pe_model = None
|
|
self.pe_tokenizer = None
|
|
self.use_pe = False
|
|
|
|
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, ErnieImageTransformer2DModel)
|
|
self.transformer = ErnieImageTransformer2DModel(
|
|
quant_config=od_config.quantization_config, **transformer_kwargs
|
|
)
|
|
|
|
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels)) if getattr(self, "vae", None) else 16
|
|
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
|
self.tokenizer_max_length = 512
|
|
self.default_sample_size = 128
|
|
|
|
self._guidance_scale = None
|
|
self._attention_kwargs = None
|
|
self._num_timesteps = None
|
|
self._current_timestep = None
|
|
self._interrupt = False
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
|
|
@staticmethod
|
|
def _distributed_prompt_sync_state() -> tuple[bool, int]:
|
|
try:
|
|
if not torch.distributed.is_available() or not torch.distributed.is_initialized():
|
|
return False, 0
|
|
if torch.distributed.get_world_size() <= 1:
|
|
return False, 0
|
|
return True, torch.distributed.get_rank()
|
|
except Exception:
|
|
return False, 0
|
|
|
|
@staticmethod
|
|
def _broadcast_enhanced_prompt(prompt: str | None) -> str | None:
|
|
values = [prompt]
|
|
torch.distributed.broadcast_object_list(values, src=0)
|
|
return values[0]
|
|
|
|
def _enhance_prompt(
|
|
self,
|
|
prompt: str,
|
|
device: torch.device,
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.95,
|
|
) -> str:
|
|
if not self.use_pe or self.pe_model is None:
|
|
return prompt
|
|
|
|
sync_prompt, rank = self._distributed_prompt_sync_state()
|
|
if sync_prompt and rank != 0:
|
|
try:
|
|
enhanced = self._broadcast_enhanced_prompt(None)
|
|
return enhanced if enhanced else prompt
|
|
except Exception as e:
|
|
logger.warning("PE enhancement broadcast failed: %s", e)
|
|
return prompt
|
|
|
|
try:
|
|
user_content = json.dumps(
|
|
{"prompt": prompt, "width": width, "height": height},
|
|
ensure_ascii=False,
|
|
)
|
|
messages = [{"role": "user", "content": user_content}]
|
|
input_text = self.pe_tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False,
|
|
)
|
|
inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(device)
|
|
output_ids = self.pe_model.generate(
|
|
**inputs,
|
|
max_new_tokens=512,
|
|
do_sample=temperature != 1.0 or top_p != 1.0,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
pad_token_id=self.pe_tokenizer.pad_token_id,
|
|
eos_token_id=self.pe_tokenizer.eos_token_id,
|
|
)
|
|
output_ids = output_ids[0][inputs.input_ids.shape[1] :]
|
|
result = self.pe_tokenizer.decode(output_ids, skip_special_tokens=True).strip()
|
|
enhanced = result if result else prompt
|
|
except Exception as e:
|
|
logger.warning("PE enhancement failed: %s", e)
|
|
enhanced = prompt
|
|
|
|
if sync_prompt:
|
|
try:
|
|
synced = self._broadcast_enhanced_prompt(enhanced)
|
|
return synced if synced else prompt
|
|
except Exception as e:
|
|
logger.warning("PE enhancement broadcast failed: %s", e)
|
|
return enhanced
|
|
|
|
@staticmethod
|
|
def _is_warmup_request(req: OmniDiffusionRequest) -> bool:
|
|
is_dummy_run = getattr(req, "is_dummy_run", None)
|
|
if callable(is_dummy_run):
|
|
return bool(is_dummy_run())
|
|
return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
|
|
|
|
@staticmethod
|
|
def _should_apply_pe(req: OmniDiffusionRequest) -> bool:
|
|
if ErnieImagePipeline._is_warmup_request(req):
|
|
return False
|
|
extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
|
|
return bool(extra_args.get("apply_pe", True))
|
|
|
|
def encode_prompt(
|
|
self,
|
|
prompt: str | list[str],
|
|
device: torch.device,
|
|
num_images_per_prompt: int = 1,
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
apply_pe: bool = True,
|
|
) -> list[torch.Tensor]:
|
|
if isinstance(prompt, str):
|
|
prompt = [prompt]
|
|
|
|
text_hiddens = []
|
|
|
|
for p in prompt:
|
|
if apply_pe and self.use_pe and self.pe_model is not None:
|
|
enhanced = self._enhance_prompt(p, device, width=width, height=height)
|
|
logger.info("PE: original='%s...' enhanced='%s...'", p[:50], enhanced[:50])
|
|
p = enhanced
|
|
ids = self.tokenizer(
|
|
p,
|
|
add_special_tokens=True,
|
|
truncation=True,
|
|
padding=False,
|
|
)["input_ids"]
|
|
|
|
if len(ids) == 0:
|
|
if self.tokenizer.bos_token_id is not None:
|
|
ids = [self.tokenizer.bos_token_id]
|
|
else:
|
|
ids = [0]
|
|
|
|
input_ids = torch.tensor([ids], device=device)
|
|
with torch.no_grad():
|
|
outputs = self.text_encoder(input_ids=input_ids, output_hidden_states=True)
|
|
hidden = outputs.hidden_states[-2][0]
|
|
|
|
for _ in range(num_images_per_prompt):
|
|
text_hiddens.append(hidden)
|
|
|
|
return text_hiddens
|
|
|
|
@staticmethod
|
|
def _patchify_latents(latents: torch.Tensor) -> torch.Tensor:
|
|
b, c, h, w = latents.shape
|
|
latents = latents.view(b, c, h // 2, 2, w // 2, 2)
|
|
latents = latents.permute(0, 1, 3, 5, 2, 4)
|
|
return latents.reshape(b, c * 4, h // 2, w // 2)
|
|
|
|
@staticmethod
|
|
def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
|
|
b, c, h, w = latents.shape
|
|
latents = latents.reshape(b, c // 4, 2, 2, h, w)
|
|
latents = latents.permute(0, 1, 4, 2, 5, 3)
|
|
return latents.reshape(b, c // 4, h * 2, w * 2)
|
|
|
|
@staticmethod
|
|
def _pad_text(text_hiddens: list[torch.Tensor], device: torch.device, dtype: torch.dtype, text_in_dim: int):
|
|
B = len(text_hiddens)
|
|
if B == 0:
|
|
return torch.zeros((0, 0, text_in_dim), device=device, dtype=dtype), torch.zeros(
|
|
(0,), device=device, dtype=torch.long
|
|
)
|
|
normalized = [
|
|
th.squeeze(1).to(device).to(dtype) if th.dim() == 3 else th.to(device).to(dtype) for th in text_hiddens
|
|
]
|
|
lens = torch.tensor([t.shape[0] for t in normalized], device=device, dtype=torch.long)
|
|
Tmax = int(lens.max().item())
|
|
text_bth = torch.zeros((B, Tmax, text_in_dim), device=device, dtype=dtype)
|
|
for i, t in enumerate(normalized):
|
|
text_bth[i, : t.shape[0], :] = t
|
|
return text_bth, lens
|
|
|
|
@property
|
|
def guidance_scale(self):
|
|
return self._guidance_scale
|
|
|
|
@property
|
|
def do_classifier_free_guidance(self):
|
|
return self._guidance_scale is not None and self._guidance_scale > 1 and not self.is_distilled
|
|
|
|
@property
|
|
def num_timesteps(self):
|
|
return self._num_timesteps
|
|
|
|
@property
|
|
def current_timestep(self):
|
|
return self._current_timestep
|
|
|
|
@property
|
|
def interrupt(self):
|
|
return self._interrupt
|
|
|
|
def _resize_dimensions(self, height: int, width: int) -> tuple[int, int]:
|
|
resized_height = height - height % self.vae_scale_factor
|
|
resized_width = width - width % self.vae_scale_factor
|
|
return resized_height, resized_width
|
|
|
|
def check_inputs(
|
|
self,
|
|
prompt,
|
|
height,
|
|
width,
|
|
prompt_embeds=None,
|
|
callback_on_step_end_tensor_inputs=None,
|
|
guidance_scale=None,
|
|
):
|
|
if (height is not None and height % self.vae_scale_factor != 0) or (
|
|
width is not None and width % self.vae_scale_factor != 0
|
|
):
|
|
resized_height, resized_width = self._resize_dimensions(height, width)
|
|
logger.warning(
|
|
"`height` and `width` have to be divisible by %s "
|
|
"but are %s and %s. Dimensions will be resized to %s and %s accordingly",
|
|
self.vae_scale_factor,
|
|
height,
|
|
width,
|
|
resized_height,
|
|
resized_width,
|
|
)
|
|
|
|
if callback_on_step_end_tensor_inputs is not None and not all(
|
|
k in ["latents", "prompt_embeds"] for k in callback_on_step_end_tensor_inputs
|
|
):
|
|
raise ValueError("`callback_on_step_end_tensor_inputs` must be a subset of ['latents', 'prompt_embeds'].")
|
|
|
|
if prompt is not None and prompt_embeds is not None:
|
|
raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please forward only one.")
|
|
if prompt is None and prompt_embeds is None:
|
|
raise ValueError("Must provide either `prompt` or `prompt_embeds`.")
|
|
|
|
def forward(
|
|
self,
|
|
req: OmniDiffusionRequest,
|
|
prompt: str | list[str] | None = None,
|
|
negative_prompt: str | list[str] | None = "",
|
|
height: int = 1024,
|
|
width: int = 1024,
|
|
num_inference_steps: int = 50,
|
|
guidance_scale: float = 4.0,
|
|
num_images_per_prompt: int = 1,
|
|
generator: torch.Generator | None = None,
|
|
latents: torch.Tensor | None = None,
|
|
prompt_embeds: list[torch.FloatTensor] | None = None,
|
|
negative_prompt_embeds: list[torch.FloatTensor] | None = None,
|
|
output_type: str = "pil",
|
|
return_dict: bool = True,
|
|
callback_on_step_end: Callable[[int, int, dict], None] | None = None,
|
|
callback_on_step_end_tensor_inputs: list[str] = ["latents"],
|
|
) -> DiffusionOutput:
|
|
if len(req.prompts) > 1:
|
|
logger.warning("This model only supports a single prompt. Taking only the first.")
|
|
first_prompt = req.prompts[0]
|
|
prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
|
|
|
|
height = req.sampling_params.height or height
|
|
width = req.sampling_params.width or width
|
|
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
|
|
guidance_scale = (
|
|
req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
|
|
)
|
|
generator = req.sampling_params.generator or generator
|
|
num_images_per_prompt = (
|
|
req.sampling_params.num_outputs_per_prompt
|
|
if req.sampling_params.num_outputs_per_prompt > 0
|
|
else num_images_per_prompt
|
|
)
|
|
|
|
if prompt is not None and isinstance(prompt, str):
|
|
batch_size = 1
|
|
elif prompt is not None and isinstance(prompt, list):
|
|
batch_size = len(prompt)
|
|
else:
|
|
batch_size = prompt_embeds.shape[0]
|
|
|
|
device = self._execution_device
|
|
dtype = self.transformer.dtype
|
|
|
|
self._guidance_scale = guidance_scale
|
|
self._current_timestep = None
|
|
self._interrupt = False
|
|
|
|
self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
|
|
height, width = self._resize_dimensions(height, width)
|
|
|
|
if prompt is not None:
|
|
if isinstance(prompt, str):
|
|
prompt = [prompt]
|
|
|
|
if prompt_embeds is not None:
|
|
text_hiddens = prompt_embeds
|
|
else:
|
|
text_hiddens = self.encode_prompt(
|
|
prompt,
|
|
device,
|
|
num_images_per_prompt,
|
|
width=width,
|
|
height=height,
|
|
apply_pe=self._should_apply_pe(req),
|
|
)
|
|
|
|
if self.do_classifier_free_guidance:
|
|
if negative_prompt_embeds is not None:
|
|
uncond_text_hiddens = negative_prompt_embeds
|
|
else:
|
|
if negative_prompt is None:
|
|
negative_prompt = ""
|
|
if isinstance(negative_prompt, str):
|
|
negative_prompt = [negative_prompt] * batch_size
|
|
uncond_text_hiddens = self.encode_prompt(
|
|
negative_prompt, device, num_images_per_prompt, width=width, height=height, apply_pe=False
|
|
)
|
|
|
|
latent_h = height // self.vae_scale_factor
|
|
latent_w = width // self.vae_scale_factor
|
|
latent_channels = self.transformer.config.in_channels
|
|
|
|
if latents is None:
|
|
latents = randn_tensor(
|
|
(batch_size * num_images_per_prompt, latent_channels, latent_h, latent_w),
|
|
generator=generator,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)
|
|
self.scheduler.set_timesteps(sigmas=sigmas[:-1], device=device)
|
|
|
|
if self.do_classifier_free_guidance:
|
|
cfg_text_hiddens = list(uncond_text_hiddens) + list(text_hiddens)
|
|
else:
|
|
cfg_text_hiddens = text_hiddens
|
|
|
|
text_bth, text_lens = self._pad_text(
|
|
text_hiddens=cfg_text_hiddens, device=device, dtype=dtype, text_in_dim=self.transformer.config.text_in_dim
|
|
)
|
|
|
|
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
|
for i, t in enumerate(self.scheduler.timesteps):
|
|
if self.interrupt:
|
|
continue
|
|
|
|
self._current_timestep = t
|
|
|
|
if self.do_classifier_free_guidance:
|
|
latent_model_input = torch.cat([latents, latents], dim=0)
|
|
t_batch = torch.full(
|
|
(batch_size * num_images_per_prompt * 2,), t.item(), device=device, dtype=dtype
|
|
)
|
|
else:
|
|
latent_model_input = latents
|
|
t_batch = torch.full((batch_size * num_images_per_prompt,), t.item(), device=device, dtype=dtype)
|
|
|
|
pred = self.transformer(
|
|
hidden_states=latent_model_input,
|
|
timestep=t_batch,
|
|
text_bth=text_bth,
|
|
text_lens=text_lens,
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
if self.do_classifier_free_guidance:
|
|
pred_uncond, pred_cond = pred.chunk(2, dim=0)
|
|
pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond)
|
|
|
|
latents = self.scheduler.step(pred, t, latents, return_dict=False)[0]
|
|
|
|
if callback_on_step_end is not None:
|
|
callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
|
|
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
|
latents = callback_outputs.pop("latents", latents)
|
|
|
|
progress_bar.update()
|
|
|
|
self._current_timestep = None
|
|
|
|
if output_type == "latent":
|
|
return DiffusionOutput(output=latents)
|
|
|
|
bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(device)
|
|
bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(device)
|
|
|
|
latents = latents * bn_std + bn_mean
|
|
|
|
latents = self._unpatchify_latents(latents)
|
|
|
|
images = self.vae.decode(latents, return_dict=False)[0]
|
|
|
|
return DiffusionOutput(
|
|
output=images, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loaded_weights = set()
|
|
transformer_weights = (
|
|
(name.replace("transformer.", "", 1), weight) for name, weight in weights if name.startswith("transformer.")
|
|
)
|
|
loaded_weights |= {f"transformer.{name}" for name in self.transformer.load_weights(transformer_weights)}
|
|
loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
|
|
loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
|
|
return loaded_weights
|