chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,438 @@
|
||||
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
|
||||
# with the following modifications:
|
||||
# - It uses the patched version of `sde_step_with_logprob` from `sd3_sde_with_logprob.py`.
|
||||
# - It returns all the intermediate latents of the denoising process as well as the log probs of each denoising step.
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers.pipelines.flux.pipeline_flux import retrieve_timesteps as retrieve_flux_timesteps
|
||||
from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
|
||||
|
||||
from .solver import run_sampling
|
||||
|
||||
|
||||
def calculate_shift(
|
||||
image_seq_len,
|
||||
base_seq_len: int = 256,
|
||||
max_seq_len: int = 4096,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
):
|
||||
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
||||
b = base_shift - m * base_seq_len
|
||||
mu = image_seq_len * m + b
|
||||
return mu
|
||||
|
||||
|
||||
def _unwrap_compiled(model):
|
||||
return model._orig_mod if hasattr(model, "_orig_mod") else model
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SD3 pipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def pipeline_with_logprob_sd3(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
prompt_3: Optional[Union[str, List[str]]] = None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 28,
|
||||
guidance_scale: float = 7.0,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
negative_prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
negative_prompt_3: Optional[Union[str, List[str]]] = None,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 256,
|
||||
noise_level: float = 0.7,
|
||||
deterministic: bool = False,
|
||||
solver: str = "flow",
|
||||
sequential_decode: bool = False,
|
||||
):
|
||||
height = height or self.default_sample_size * self.vae_scale_factor
|
||||
width = width or self.default_sample_size * self.vae_scale_factor
|
||||
|
||||
self.check_inputs(
|
||||
prompt,
|
||||
prompt_2,
|
||||
prompt_3,
|
||||
height,
|
||||
width,
|
||||
negative_prompt=negative_prompt,
|
||||
negative_prompt_2=negative_prompt_2,
|
||||
negative_prompt_3=negative_prompt_3,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length=max_sequence_length,
|
||||
)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
self._joint_attention_kwargs = joint_attention_kwargs
|
||||
self._current_timestep = None
|
||||
self._interrupt = False
|
||||
|
||||
# 2. Define call parameters
|
||||
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
|
||||
|
||||
lora_scale = self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
|
||||
(
|
||||
prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
) = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_3=prompt_3,
|
||||
negative_prompt=negative_prompt,
|
||||
negative_prompt_2=negative_prompt_2,
|
||||
negative_prompt_3=negative_prompt_3,
|
||||
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
if self.do_classifier_free_guidance:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
|
||||
|
||||
# 4. Prepare latent variables
|
||||
num_channels_latents = self.transformer.config.in_channels
|
||||
if latents is None:
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=None,
|
||||
)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
sigmas = self.scheduler.sigmas.float()
|
||||
|
||||
def v_pred_fn(z, sigma):
|
||||
latent_model_input = torch.cat([z] * 2) if self.do_classifier_free_guidance else z
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timesteps = torch.full([latent_model_input.shape[0]], sigma * 1000, device=z.device, dtype=torch.long)
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
timestep=timesteps,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
pooled_projections=pooled_prompt_embeds,
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
noise_pred = noise_pred.to(prompt_embeds.dtype)
|
||||
if self.do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
return noise_pred
|
||||
|
||||
# 6. Prepare image embeddings
|
||||
all_latents = [latents]
|
||||
all_log_probs = []
|
||||
|
||||
# 7. Denoising loop
|
||||
latents, all_latents, all_log_probs = run_sampling(v_pred_fn, latents, sigmas, solver, deterministic, noise_level)
|
||||
|
||||
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
|
||||
latents = latents.to(dtype=self.vae.dtype)
|
||||
if sequential_decode and latents.shape[0] > 1:
|
||||
decoded_batches = []
|
||||
for idx in range(latents.shape[0]):
|
||||
decoded_batches.append(self.vae.decode(latents[idx : idx + 1], return_dict=False)[0])
|
||||
image = torch.cat(decoded_batches, dim=0)
|
||||
else:
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
return image, all_latents, all_log_probs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FLUX.1 pipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def pipeline_with_logprob_flux(
|
||||
pipeline,
|
||||
prompt=None,
|
||||
prompt_2=None,
|
||||
height=None,
|
||||
width=None,
|
||||
num_inference_steps=28,
|
||||
guidance_scale=3.5,
|
||||
num_images_per_prompt=1,
|
||||
generator=None,
|
||||
latents=None,
|
||||
prompt_embeds=None,
|
||||
pooled_prompt_embeds=None,
|
||||
text_ids=None,
|
||||
output_type="pt",
|
||||
joint_attention_kwargs=None,
|
||||
max_sequence_length=512,
|
||||
noise_level=0.7,
|
||||
deterministic=False,
|
||||
solver="flow",
|
||||
sequential_decode=False,
|
||||
):
|
||||
height = height or pipeline.default_sample_size * pipeline.vae_scale_factor
|
||||
width = width or pipeline.default_sample_size * pipeline.vae_scale_factor
|
||||
|
||||
pipeline.check_inputs(
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
max_sequence_length=max_sequence_length,
|
||||
)
|
||||
|
||||
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 = pipeline._execution_device
|
||||
lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
|
||||
|
||||
if prompt_embeds is None or pooled_prompt_embeds is None or text_ids is None:
|
||||
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
num_channels_latents = pipeline.transformer.config.in_channels // 4
|
||||
if latents is None:
|
||||
latents, latent_image_ids = pipeline.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
latent_image_ids = pipeline._prepare_latent_image_ids(
|
||||
batch_size * num_images_per_prompt,
|
||||
height // pipeline.vae_scale_factor,
|
||||
width // pipeline.vae_scale_factor,
|
||||
device,
|
||||
prompt_embeds.dtype,
|
||||
)
|
||||
|
||||
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
|
||||
if hasattr(pipeline.scheduler.config, "use_flow_sigmas") and pipeline.scheduler.config.use_flow_sigmas:
|
||||
sigmas = None
|
||||
|
||||
image_seq_len = latents.shape[1]
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
pipeline.scheduler.config.get("base_image_seq_len", 256),
|
||||
pipeline.scheduler.config.get("max_image_seq_len", 4096),
|
||||
pipeline.scheduler.config.get("base_shift", 0.5),
|
||||
pipeline.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
_, num_inference_steps = retrieve_flux_timesteps(
|
||||
pipeline.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
mu=mu,
|
||||
)
|
||||
sigmas = pipeline.scheduler.sigmas.float()
|
||||
|
||||
active_transformer = pipeline.transformer
|
||||
guidance_config = _unwrap_compiled(active_transformer).config
|
||||
if guidance_config.guidance_embeds:
|
||||
guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0])
|
||||
else:
|
||||
guidance = None
|
||||
|
||||
def v_pred_fn(z, sigma):
|
||||
timestep = torch.full([z.shape[0]], float(sigma), device=z.device, dtype=z.dtype)
|
||||
noise_pred = active_transformer(
|
||||
hidden_states=z,
|
||||
timestep=timestep,
|
||||
guidance=guidance,
|
||||
pooled_projections=pooled_prompt_embeds,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=latent_image_ids,
|
||||
joint_attention_kwargs=joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
return noise_pred
|
||||
|
||||
all_latents = [latents]
|
||||
latents, all_latents, all_log_probs = run_sampling(v_pred_fn, latents, sigmas, solver, deterministic, noise_level)
|
||||
|
||||
latents = pipeline._unpack_latents(latents, height, width, pipeline.vae_scale_factor)
|
||||
latents = (latents / pipeline.vae.config.scaling_factor) + pipeline.vae.config.shift_factor
|
||||
latents = latents.to(dtype=pipeline.vae.dtype)
|
||||
|
||||
if sequential_decode and latents.shape[0] > 1:
|
||||
decoded_batches = []
|
||||
for idx in range(latents.shape[0]):
|
||||
decoded_batches.append(pipeline.vae.decode(latents[idx : idx + 1], return_dict=False)[0])
|
||||
image = torch.cat(decoded_batches, dim=0)
|
||||
else:
|
||||
image = pipeline.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
image = pipeline.image_processor.postprocess(image, output_type=output_type)
|
||||
pipeline.maybe_free_model_hooks()
|
||||
|
||||
return image, all_latents, latent_image_ids, text_ids, all_log_probs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sana pipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def pipeline_with_logprob_sana(
|
||||
transformer,
|
||||
vae,
|
||||
*,
|
||||
latents=None,
|
||||
num_channels=None,
|
||||
latent_size=None,
|
||||
prompt_embeds=None,
|
||||
prompt_attention_mask=None,
|
||||
negative_prompt_embeds=None,
|
||||
negative_prompt_attention_mask=None,
|
||||
num_inference_steps=20,
|
||||
guidance_scale=4.5,
|
||||
noise_level=0.7,
|
||||
deterministic=False,
|
||||
sequential_decode=False,
|
||||
solver="flow",
|
||||
):
|
||||
assert prompt_embeds is not None
|
||||
|
||||
if latents is None:
|
||||
assert num_channels is not None and latent_size is not None
|
||||
latents = torch.randn(
|
||||
prompt_embeds.shape[0],
|
||||
num_channels,
|
||||
latent_size,
|
||||
latent_size,
|
||||
device=prompt_embeds.device,
|
||||
dtype=prompt_embeds.dtype,
|
||||
)
|
||||
|
||||
device = latents.device
|
||||
dtype = latents.dtype
|
||||
sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device, dtype=dtype)
|
||||
|
||||
do_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
|
||||
|
||||
caption_4d = prompt_embeds.unsqueeze(1) if prompt_embeds.dim() == 3 else prompt_embeds
|
||||
mask_4d = (
|
||||
prompt_attention_mask.unsqueeze(1).unsqueeze(1).to(torch.int16)
|
||||
if prompt_attention_mask is not None and prompt_attention_mask.dim() == 2
|
||||
else prompt_attention_mask
|
||||
)
|
||||
|
||||
if do_cfg:
|
||||
neg_4d = negative_prompt_embeds.unsqueeze(1) if negative_prompt_embeds.dim() == 3 else negative_prompt_embeds
|
||||
neg_mask_4d = (
|
||||
negative_prompt_attention_mask.unsqueeze(1).unsqueeze(1).to(torch.int16)
|
||||
if negative_prompt_attention_mask is not None and negative_prompt_attention_mask.dim() == 2
|
||||
else negative_prompt_attention_mask
|
||||
)
|
||||
y_in = torch.cat([neg_4d, caption_4d], dim=0)
|
||||
m_in = torch.cat([neg_mask_4d, mask_4d], dim=0) if mask_4d is not None else None
|
||||
else:
|
||||
y_in = caption_4d
|
||||
m_in = mask_4d
|
||||
|
||||
def v_pred_fn(z, sigma):
|
||||
z_in = torch.cat([z, z], dim=0) if do_cfg else z
|
||||
t_batch = sigma.expand(z_in.shape[0]).to(device)
|
||||
pred = transformer(z_in, t_batch, y_in, mask=m_in)
|
||||
if do_cfg:
|
||||
u, c = pred.chunk(2)
|
||||
pred = u + guidance_scale * (c - u)
|
||||
return pred
|
||||
|
||||
latents, all_latents, _ = run_sampling(
|
||||
v_pred_fn,
|
||||
latents,
|
||||
sigmas,
|
||||
solver,
|
||||
deterministic,
|
||||
noise_level,
|
||||
)
|
||||
|
||||
vae_dtype = next(vae.parameters()).dtype
|
||||
latents_dec = latents.to(vae_dtype) / vae.config.scaling_factor
|
||||
if sequential_decode and latents_dec.shape[0] > 1:
|
||||
decoded = []
|
||||
for idx in range(latents_dec.shape[0]):
|
||||
decoded.append(vae.decode(latents_dec[idx : idx + 1], return_dict=False)[0])
|
||||
image = torch.cat(decoded, dim=0)
|
||||
else:
|
||||
image = vae.decode(latents_dec, return_dict=False)[0]
|
||||
images = (image / 2 + 0.5).clamp(0, 1)
|
||||
|
||||
return images, all_latents, sigmas[:-1]
|
||||
@@ -0,0 +1,357 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import tqdm
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
tqdm = partial(tqdm.tqdm, dynamic_ncols=True)
|
||||
|
||||
|
||||
# Modified from MixGRPO
|
||||
def run_sampling(
|
||||
v_pred_fn,
|
||||
z,
|
||||
sigma_schedule,
|
||||
solver="flow",
|
||||
determistic=False,
|
||||
eta=0.7,
|
||||
):
|
||||
assert solver in ["flow", "dance", "ddim", "dpm1", "dpm2"]
|
||||
dtype = z.dtype
|
||||
all_latents = [z]
|
||||
all_log_probs = []
|
||||
|
||||
if "dpm" in solver:
|
||||
order = int(solver[-1])
|
||||
dpm_state = DPMState(order=order)
|
||||
for i in tqdm(
|
||||
range(len(sigma_schedule) - 1),
|
||||
desc="Sampling Progress",
|
||||
disable=not dist.is_initialized() or dist.get_rank() != 0,
|
||||
):
|
||||
sigma = sigma_schedule[i]
|
||||
|
||||
pred = v_pred_fn(z.to(dtype), sigma)
|
||||
if solver == "flow":
|
||||
z, pred_original, log_prob = flow_grpo_step(
|
||||
model_output=pred.float(),
|
||||
latents=z.float(),
|
||||
eta=eta if not determistic else 0,
|
||||
sigmas=sigma_schedule,
|
||||
index=i,
|
||||
prev_sample=None,
|
||||
)
|
||||
elif solver == "dance":
|
||||
z, pred_original, log_prob = dance_grpo_step(
|
||||
pred.float(), z.float(), eta if not determistic else 0, sigmas=sigma_schedule, index=i, prev_sample=None
|
||||
)
|
||||
elif solver == "ddim":
|
||||
z, pred_original, log_prob = ddim_step(
|
||||
pred.float(), z.float(), eta if not determistic else 0, sigmas=sigma_schedule, index=i, prev_sample=None
|
||||
)
|
||||
elif "dpm" in solver:
|
||||
assert determistic
|
||||
z, pred_original, log_prob = dpm_step(
|
||||
order,
|
||||
model_output=pred.float(),
|
||||
sample=z.float(),
|
||||
step_index=i,
|
||||
timesteps=sigma_schedule[:-1],
|
||||
sigmas=sigma_schedule,
|
||||
dpm_state=dpm_state,
|
||||
)
|
||||
else:
|
||||
assert False
|
||||
z = z.to(dtype)
|
||||
all_latents.append(z)
|
||||
all_log_probs.append(log_prob)
|
||||
|
||||
latents = z.to(dtype)
|
||||
# all_latents = torch.stack(all_latents, dim=1) # (batch_size, num_steps + 1, 4, 64, 64)
|
||||
# all_log_probs = torch.stack(all_log_probs, dim=1) # (batch_size, num_steps, 1)
|
||||
return latents, all_latents, all_log_probs
|
||||
|
||||
|
||||
def flow_grpo_step(
|
||||
model_output: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
eta: float,
|
||||
sigmas: torch.Tensor,
|
||||
index: int,
|
||||
prev_sample: torch.Tensor,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
):
|
||||
device = model_output.device
|
||||
sigma = sigmas[index].to(device)
|
||||
sigma_prev = sigmas[index + 1].to(device)
|
||||
sigma_max = sigmas[1].item()
|
||||
dt = sigma_prev - sigma # neg dt
|
||||
|
||||
pred_original_sample = latents - sigma * model_output
|
||||
|
||||
std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * eta
|
||||
|
||||
if prev_sample is not None and generator is not None:
|
||||
raise ValueError(
|
||||
"Cannot pass both generator and prev_sample. Please make sure that either `generator` or"
|
||||
" `prev_sample` stays `None`."
|
||||
)
|
||||
|
||||
prev_sample_mean = (
|
||||
latents * (1 + std_dev_t**2 / (2 * sigma) * dt)
|
||||
+ model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
|
||||
)
|
||||
|
||||
if prev_sample is None:
|
||||
variance_noise = randn_tensor(model_output.shape, generator=generator, device=device, dtype=model_output.dtype)
|
||||
prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
|
||||
|
||||
log_prob = (
|
||||
-((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
|
||||
- torch.log(std_dev_t * torch.sqrt(-1 * dt))
|
||||
- torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
|
||||
)
|
||||
|
||||
# mean along all but batch dimension
|
||||
log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
|
||||
|
||||
return prev_sample, pred_original_sample, log_prob
|
||||
|
||||
|
||||
def dance_grpo_step(
|
||||
model_output: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
eta: float,
|
||||
sigmas: torch.Tensor,
|
||||
index: int,
|
||||
prev_sample: torch.Tensor,
|
||||
):
|
||||
sigma = sigmas[index]
|
||||
dsigma = sigmas[index + 1] - sigma # neg dt
|
||||
prev_sample_mean = latents + dsigma * model_output
|
||||
|
||||
pred_original_sample = latents - sigma * model_output
|
||||
|
||||
delta_t = sigma - sigmas[index + 1] # pos -dt
|
||||
std_dev_t = eta * math.sqrt(delta_t)
|
||||
|
||||
score_estimate = -(latents - pred_original_sample * (1 - sigma)) / sigma**2
|
||||
log_term = -0.5 * eta**2 * score_estimate
|
||||
prev_sample_mean = prev_sample_mean + log_term * dsigma
|
||||
|
||||
if prev_sample is None:
|
||||
prev_sample = prev_sample_mean + torch.randn_like(prev_sample_mean) * std_dev_t
|
||||
|
||||
# log prob of prev_sample given prev_sample_mean and std_dev_t
|
||||
log_prob = -((prev_sample.detach().to(torch.float32) - prev_sample_mean.to(torch.float32)) ** 2) / (
|
||||
2 * (std_dev_t**2)
|
||||
)
|
||||
-math.log(std_dev_t) - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
|
||||
|
||||
# mean along all but batch dimension
|
||||
log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
|
||||
return prev_sample, pred_original_sample, log_prob
|
||||
|
||||
|
||||
def ddim_step(
|
||||
model_output: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
eta: float,
|
||||
sigmas: torch.Tensor,
|
||||
index: int,
|
||||
prev_sample: torch.Tensor,
|
||||
):
|
||||
model_output = convert_model_output(model_output, latents, sigmas, step_index=index)
|
||||
prev_sample, prev_sample_mean, std_dev_t, dt_sqrt = ddim_update(
|
||||
model_output,
|
||||
sigmas.to(torch.float64),
|
||||
index,
|
||||
latents,
|
||||
eta=eta,
|
||||
)
|
||||
|
||||
# Compute log_prob
|
||||
log_prob = (
|
||||
-((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * dt_sqrt) ** 2))
|
||||
- torch.log(std_dev_t * dt_sqrt)
|
||||
- torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
|
||||
)
|
||||
|
||||
# mean along all but batch dimension
|
||||
log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
|
||||
return prev_sample, model_output, log_prob
|
||||
|
||||
|
||||
@dataclass
|
||||
class DPMState:
|
||||
order: int
|
||||
model_outputs: List[torch.Tensor] = None
|
||||
lower_order_nums = 0
|
||||
|
||||
def __post_init__(self):
|
||||
self.model_outputs = [None] * self.order
|
||||
|
||||
def update(self, model_output: torch.Tensor):
|
||||
for i in range(self.order - 1):
|
||||
self.model_outputs[i] = self.model_outputs[i + 1]
|
||||
self.model_outputs[-1] = model_output
|
||||
|
||||
def update_lower_order(self):
|
||||
if self.lower_order_nums < self.order:
|
||||
self.lower_order_nums += 1
|
||||
|
||||
|
||||
def dpm_step(
|
||||
order,
|
||||
model_output: torch.Tensor,
|
||||
sample: torch.Tensor,
|
||||
step_index: int,
|
||||
timesteps: list,
|
||||
sigmas: torch.Tensor,
|
||||
dpm_state: DPMState = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
# Improve numerical stability for small number of steps
|
||||
lower_order_final = step_index == len(timesteps) - 1
|
||||
lower_order_second = (step_index == len(timesteps) - 2) and len(timesteps) < 15
|
||||
|
||||
model_output = convert_model_output(model_output, sample, sigmas, step_index=step_index)
|
||||
|
||||
assert dpm_state is not None
|
||||
dpm_state.update(model_output)
|
||||
|
||||
# Upcast to avoid precision issues when computing prev_sample
|
||||
sample = sample.to(torch.float32)
|
||||
|
||||
if order == 1 or dpm_state.lower_order_nums < 1 or lower_order_final:
|
||||
if step_index == 0 or lower_order_final:
|
||||
prev_sample, _, _, _ = ddim_update(
|
||||
model_output,
|
||||
sigmas.to(torch.float64),
|
||||
step_index,
|
||||
sample,
|
||||
eta=0.0,
|
||||
)
|
||||
else:
|
||||
prev_sample = dpm_solver_first_order_update(
|
||||
model_output,
|
||||
sigmas.to(torch.float64),
|
||||
step_index,
|
||||
sample,
|
||||
)
|
||||
elif order == 2 or dpm_state.lower_order_nums < 2 or lower_order_second:
|
||||
prev_sample = multistep_dpm_solver_second_order_update(
|
||||
dpm_state.model_outputs,
|
||||
sigmas.to(torch.float64),
|
||||
step_index,
|
||||
sample,
|
||||
)
|
||||
else:
|
||||
assert False
|
||||
|
||||
dpm_state.update_lower_order()
|
||||
|
||||
# Cast sample back to expected dtype
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
return prev_sample, model_output, None
|
||||
|
||||
|
||||
def convert_model_output(
|
||||
model_output,
|
||||
sample,
|
||||
sigmas,
|
||||
step_index,
|
||||
) -> torch.Tensor:
|
||||
sigma_t = sigmas[step_index]
|
||||
x0_pred = sample - sigma_t * model_output
|
||||
|
||||
return x0_pred
|
||||
|
||||
|
||||
def ddim_update(
|
||||
model_output: torch.Tensor,
|
||||
sigmas,
|
||||
step_index,
|
||||
sample: torch.Tensor = None,
|
||||
noise: Optional[torch.Tensor] = None,
|
||||
eta: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
|
||||
t, s = sigmas[step_index + 1], sigmas[step_index]
|
||||
|
||||
std_dev_t = eta * t
|
||||
dt_sqrt = torch.sqrt(1.0 - t**2 * (1 - s) ** 2 / (s**2 * (1 - t) ** 2))
|
||||
rho_t = std_dev_t * dt_sqrt
|
||||
noise_pred = (sample - (1 - s) * model_output) / s
|
||||
if noise is None:
|
||||
noise = torch.randn_like(model_output)
|
||||
prev_mean = (1 - t) * model_output + torch.sqrt(t**2 - rho_t**2) * noise_pred
|
||||
x_t = prev_mean + rho_t * noise
|
||||
|
||||
return x_t, prev_mean, std_dev_t, dt_sqrt
|
||||
|
||||
|
||||
def dpm_solver_first_order_update(
|
||||
model_output: torch.Tensor,
|
||||
sigmas,
|
||||
step_index,
|
||||
sample: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
sigma_t, sigma_s = sigmas[step_index + 1], sigmas[step_index]
|
||||
alpha_t, sigma_t = _sigma_to_alpha_sigma_t(sigma_t)
|
||||
alpha_s, sigma_s = _sigma_to_alpha_sigma_t(sigma_s)
|
||||
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
||||
lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
|
||||
|
||||
h = lambda_t - lambda_s
|
||||
x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
|
||||
|
||||
return x_t
|
||||
|
||||
|
||||
def multistep_dpm_solver_second_order_update(
|
||||
model_output_list: List[torch.Tensor],
|
||||
sigmas,
|
||||
step_index,
|
||||
sample: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
sigma_t, sigma_s0, sigma_s1 = (
|
||||
sigmas[step_index + 1],
|
||||
sigmas[step_index],
|
||||
sigmas[step_index - 1],
|
||||
)
|
||||
|
||||
alpha_t, sigma_t = _sigma_to_alpha_sigma_t(sigma_t)
|
||||
alpha_s0, sigma_s0 = _sigma_to_alpha_sigma_t(sigma_s0)
|
||||
alpha_s1, sigma_s1 = _sigma_to_alpha_sigma_t(sigma_s1)
|
||||
|
||||
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
||||
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
|
||||
lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
|
||||
|
||||
m0, m1 = model_output_list[-1], model_output_list[-2]
|
||||
|
||||
h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
|
||||
r0 = h_0 / h
|
||||
D0, D1 = m0, (1.0 / r0) * (m0 - m1)
|
||||
|
||||
x_t = (
|
||||
(sigma_t / sigma_s0) * sample
|
||||
- (alpha_t * (torch.exp(-h) - 1.0)) * D0
|
||||
- 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
|
||||
)
|
||||
|
||||
return x_t
|
||||
|
||||
|
||||
def _sigma_to_alpha_sigma_t(sigma):
|
||||
alpha_t = 1 - sigma
|
||||
sigma_t = sigma
|
||||
return alpha_t, sigma_t
|
||||
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _as_prompt_list(prompt):
|
||||
return [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
|
||||
def _move_to_device(value, device):
|
||||
if value is None or device is None:
|
||||
return value
|
||||
return value.to(device)
|
||||
|
||||
|
||||
def _encode_prompt_with_t5(
|
||||
text_encoder,
|
||||
tokenizer,
|
||||
max_sequence_length,
|
||||
prompt=None,
|
||||
num_images_per_prompt=1,
|
||||
device=None,
|
||||
text_input_ids=None,
|
||||
):
|
||||
prompt = _as_prompt_list(prompt)
|
||||
batch_size = len(prompt)
|
||||
|
||||
if tokenizer is not None:
|
||||
text_inputs = tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
elif text_input_ids is None:
|
||||
raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
|
||||
|
||||
prompt_embeds = text_encoder(text_input_ids.to(device))[0]
|
||||
prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
|
||||
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
||||
return prompt_embeds
|
||||
|
||||
|
||||
def _encode_prompt_with_clip(
|
||||
text_encoder,
|
||||
tokenizer,
|
||||
prompt,
|
||||
device=None,
|
||||
text_input_ids=None,
|
||||
num_images_per_prompt=1,
|
||||
):
|
||||
prompt = _as_prompt_list(prompt)
|
||||
batch_size = len(prompt)
|
||||
|
||||
if tokenizer is not None:
|
||||
text_inputs = tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=77,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
elif text_input_ids is None:
|
||||
raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
|
||||
|
||||
prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
|
||||
pooled_prompt_embeds = prompt_embeds[0]
|
||||
prompt_embeds = prompt_embeds.hidden_states[-2]
|
||||
prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
|
||||
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
||||
return prompt_embeds, pooled_prompt_embeds
|
||||
|
||||
|
||||
def encode_sd3_prompt(
|
||||
text_encoders,
|
||||
tokenizers,
|
||||
prompt,
|
||||
max_sequence_length,
|
||||
device=None,
|
||||
num_images_per_prompt=1,
|
||||
text_input_ids_list=None,
|
||||
):
|
||||
prompt = _as_prompt_list(prompt)
|
||||
clip_prompt_embeds_list = []
|
||||
clip_pooled_prompt_embeds_list = []
|
||||
|
||||
for idx, (tokenizer, text_encoder) in enumerate(zip(tokenizers[:2], text_encoders[:2])):
|
||||
prompt_embeds, pooled_prompt_embeds = _encode_prompt_with_clip(
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
prompt=prompt,
|
||||
device=device if device is not None else text_encoder.device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
text_input_ids=text_input_ids_list[idx] if text_input_ids_list else None,
|
||||
)
|
||||
clip_prompt_embeds_list.append(prompt_embeds)
|
||||
clip_pooled_prompt_embeds_list.append(pooled_prompt_embeds)
|
||||
|
||||
clip_prompt_embeds = torch.cat(clip_prompt_embeds_list, dim=-1)
|
||||
pooled_prompt_embeds = torch.cat(clip_pooled_prompt_embeds_list, dim=-1)
|
||||
|
||||
t5_prompt_embed = _encode_prompt_with_t5(
|
||||
text_encoders[-1],
|
||||
tokenizers[-1],
|
||||
max_sequence_length,
|
||||
prompt=prompt,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
text_input_ids=text_input_ids_list[-1] if text_input_ids_list else None,
|
||||
device=device if device is not None else text_encoders[-1].device,
|
||||
)
|
||||
|
||||
clip_prompt_embeds = torch.nn.functional.pad(
|
||||
clip_prompt_embeds,
|
||||
(0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]),
|
||||
)
|
||||
prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
|
||||
|
||||
target_device = device if device is not None else prompt_embeds.device
|
||||
return _move_to_device(prompt_embeds, target_device), _move_to_device(pooled_prompt_embeds, target_device)
|
||||
|
||||
|
||||
def encode_flux_prompt(
|
||||
pipeline,
|
||||
prompt,
|
||||
max_sequence_length,
|
||||
device=None,
|
||||
num_images_per_prompt=1,
|
||||
prompt_2=None,
|
||||
lora_scale=None,
|
||||
):
|
||||
prompt = _as_prompt_list(prompt)
|
||||
prompt_2 = prompt if prompt_2 is None else _as_prompt_list(prompt_2)
|
||||
|
||||
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_embeds=None,
|
||||
pooled_prompt_embeds=None,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
target_device = device if device is not None else prompt_embeds.device
|
||||
return (
|
||||
_move_to_device(prompt_embeds, target_device),
|
||||
_move_to_device(pooled_prompt_embeds, target_device),
|
||||
_move_to_device(text_ids, target_device),
|
||||
)
|
||||
|
||||
|
||||
def encode_sana_prompt(
|
||||
pipeline,
|
||||
prompt,
|
||||
max_sequence_length,
|
||||
device=None,
|
||||
negative_prompt="",
|
||||
do_classifier_free_guidance=True,
|
||||
):
|
||||
prompt = _as_prompt_list(prompt)
|
||||
prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = (
|
||||
pipeline.encode_prompt(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
device=device,
|
||||
max_sequence_length=max_sequence_length,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
)
|
||||
)
|
||||
|
||||
target_device = device if device is not None else prompt_embeds.device
|
||||
return (
|
||||
_move_to_device(prompt_embeds, target_device),
|
||||
_move_to_device(prompt_attention_mask, target_device),
|
||||
_move_to_device(negative_prompt_embeds, target_device),
|
||||
_move_to_device(negative_prompt_attention_mask, target_device),
|
||||
)
|
||||
Reference in New Issue
Block a user