Spaces:
Runtime error
Runtime error
File size: 4,118 Bytes
6b803ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
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
|