|
import torch |
|
import math |
|
|
|
from diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc |
|
from diffusers_helper.k_diffusion.wrapper import fm_wrapper |
|
from diffusers_helper.utils import repeat_to_batch_size |
|
|
|
|
|
def flux_time_shift(t, mu=1.15, sigma=1.0): |
|
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) |
|
|
|
|
|
def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, exp_max=7.0): |
|
k = (y2 - y1) / (x2 - x1) |
|
b = y1 - k * x1 |
|
mu = k * context_length + b |
|
mu = min(mu, math.log(exp_max)) |
|
return mu |
|
|
|
|
|
def get_flux_sigmas_from_mu(n, mu): |
|
sigmas = torch.linspace(1, 0, steps=n + 1) |
|
sigmas = flux_time_shift(sigmas, mu=mu) |
|
return sigmas |
|
|
|
|
|
@torch.inference_mode() |
|
def sample_hunyuan( |
|
transformer, |
|
sampler='unipc', |
|
initial_latent=None, |
|
concat_latent=None, |
|
strength=1.0, |
|
width=512, |
|
height=512, |
|
frames=16, |
|
real_guidance_scale=1.0, |
|
distilled_guidance_scale=6.0, |
|
guidance_rescale=0.0, |
|
shift=None, |
|
num_inference_steps=25, |
|
batch_size=None, |
|
generator=None, |
|
prompt_embeds=None, |
|
prompt_embeds_mask=None, |
|
prompt_poolers=None, |
|
negative_prompt_embeds=None, |
|
negative_prompt_embeds_mask=None, |
|
negative_prompt_poolers=None, |
|
dtype=torch.bfloat16, |
|
device=None, |
|
negative_kwargs=None, |
|
callback=None, |
|
**kwargs, |
|
): |
|
device = device or transformer.device |
|
|
|
if batch_size is None: |
|
batch_size = int(prompt_embeds.shape[0]) |
|
|
|
latents = torch.randn((batch_size, 16, (frames + 3) // 4, height // 8, width // 8), generator=generator, device=generator.device).to(device=device, dtype=torch.float32) |
|
|
|
B, C, T, H, W = latents.shape |
|
seq_length = T * H * W // 4 |
|
|
|
if shift is None: |
|
mu = calculate_flux_mu(seq_length, exp_max=7.0) |
|
else: |
|
mu = math.log(shift) |
|
|
|
sigmas = get_flux_sigmas_from_mu(num_inference_steps, mu).to(device) |
|
|
|
k_model = fm_wrapper(transformer) |
|
|
|
if initial_latent is not None: |
|
sigmas = sigmas * strength |
|
first_sigma = sigmas[0].to(device=device, dtype=torch.float32) |
|
initial_latent = initial_latent.to(device=device, dtype=torch.float32) |
|
latents = initial_latent.float() * (1.0 - first_sigma) + latents.float() * first_sigma |
|
|
|
if concat_latent is not None: |
|
concat_latent = concat_latent.to(latents) |
|
|
|
distilled_guidance = torch.tensor([distilled_guidance_scale * 1000.0] * batch_size).to(device=device, dtype=dtype) |
|
|
|
prompt_embeds = repeat_to_batch_size(prompt_embeds, batch_size) |
|
prompt_embeds_mask = repeat_to_batch_size(prompt_embeds_mask, batch_size) |
|
prompt_poolers = repeat_to_batch_size(prompt_poolers, batch_size) |
|
negative_prompt_embeds = repeat_to_batch_size(negative_prompt_embeds, batch_size) |
|
negative_prompt_embeds_mask = repeat_to_batch_size(negative_prompt_embeds_mask, batch_size) |
|
negative_prompt_poolers = repeat_to_batch_size(negative_prompt_poolers, batch_size) |
|
concat_latent = repeat_to_batch_size(concat_latent, batch_size) |
|
|
|
sampler_kwargs = dict( |
|
dtype=dtype, |
|
cfg_scale=real_guidance_scale, |
|
cfg_rescale=guidance_rescale, |
|
concat_latent=concat_latent, |
|
positive=dict( |
|
pooled_projections=prompt_poolers, |
|
encoder_hidden_states=prompt_embeds, |
|
encoder_attention_mask=prompt_embeds_mask, |
|
guidance=distilled_guidance, |
|
**kwargs, |
|
), |
|
negative=dict( |
|
pooled_projections=negative_prompt_poolers, |
|
encoder_hidden_states=negative_prompt_embeds, |
|
encoder_attention_mask=negative_prompt_embeds_mask, |
|
guidance=distilled_guidance, |
|
**(kwargs if negative_kwargs is None else {**kwargs, **negative_kwargs}), |
|
) |
|
) |
|
|
|
if sampler == 'unipc': |
|
results = sample_unipc(k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False, callback=callback) |
|
else: |
|
raise NotImplementedError(f'Sampler {sampler} is not supported.') |
|
|
|
return results |
|
|
