├── requirements.txt
├── diffusers_helper
├── hf_login.py
├── clip_vision.py
├── bucket_tools.py
├── dit_common.py
├── k_diffusion
│ ├── wrapper.py
│ └── uni_pc_fm.py
├── thread_utils.py
├── gradio
│ └── progress_bar.py
├── hunyuan.py
├── pipelines
│ └── k_diffusion_hunyuan.py
├── memory.py
├── utils.py
└── models
│ └── hunyuan_video_packed.py
├── .gitignore
├── utils
└── lora_utils.py
├── LICENSE
├── README.md
├── demo_gradio_f1.py
└── demo_gradio.py
/requirements.txt:
--------------------------------------------------------------------------------
1 | accelerate==1.6.0
2 | diffusers==0.33.1
3 | transformers==4.46.2
4 | gradio==5.23.2
5 | sentencepiece==0.2.0
6 | pillow==11.1.0
7 | av==12.1.0
8 | numpy==1.26.2
9 | scipy==1.12.0
10 | requests==2.31.0
11 | torchsde==0.2.6
12 |
13 | einops
14 | opencv-contrib-python
15 | safetensors
16 |
--------------------------------------------------------------------------------
/diffusers_helper/hf_login.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 |
4 | def login(token):
5 | from huggingface_hub import login
6 | import time
7 |
8 | while True:
9 | try:
10 | login(token)
11 | print('HF login ok.')
12 | break
13 | except Exception as e:
14 | print(f'HF login failed: {e}. Retrying')
15 | time.sleep(0.5)
16 |
17 |
18 | hf_token = os.environ.get('HF_TOKEN', None)
19 |
20 | if hf_token is not None:
21 | login(hf_token)
22 |
--------------------------------------------------------------------------------
/diffusers_helper/clip_vision.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | def hf_clip_vision_encode(image, feature_extractor, image_encoder):
5 | assert isinstance(image, np.ndarray)
6 | assert image.ndim == 3 and image.shape[2] == 3
7 | assert image.dtype == np.uint8
8 |
9 | preprocessed = feature_extractor.preprocess(images=image, return_tensors="pt").to(device=image_encoder.device, dtype=image_encoder.dtype)
10 | image_encoder_output = image_encoder(**preprocessed)
11 |
12 | return image_encoder_output
13 |
--------------------------------------------------------------------------------
/diffusers_helper/bucket_tools.py:
--------------------------------------------------------------------------------
1 | bucket_options = {
2 | (416, 960),
3 | (448, 864),
4 | (480, 832),
5 | (512, 768),
6 | (544, 704),
7 | (576, 672),
8 | (608, 640),
9 | (640, 608),
10 | (672, 576),
11 | (704, 544),
12 | (768, 512),
13 | (832, 480),
14 | (864, 448),
15 | (960, 416),
16 | }
17 |
18 |
19 | def find_nearest_bucket(h, w, resolution=640):
20 | min_metric = float('inf')
21 | best_bucket = None
22 | for (bucket_h, bucket_w) in bucket_options:
23 | metric = abs(h * bucket_w - w * bucket_h)
24 | if metric <= min_metric:
25 | min_metric = metric
26 | best_bucket = (bucket_h, bucket_w)
27 |
28 | if resolution != 640:
29 | scale_factor = resolution / 640.0
30 | scaled_height = round(best_bucket[0] * scale_factor / 16) * 16
31 | scaled_width = round(best_bucket[1] * scale_factor / 16) * 16
32 | best_bucket = (scaled_height, scaled_width)
33 | print(f'Resolution: {best_bucket[1]} x {best_bucket[0]}')
34 |
35 | return best_bucket
36 |
37 |
--------------------------------------------------------------------------------
/diffusers_helper/dit_common.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import accelerate.accelerator
3 |
4 | from diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm, AdaLayerNormContinuous
5 |
6 |
7 | accelerate.accelerator.convert_outputs_to_fp32 = lambda x: x
8 |
9 |
10 | def LayerNorm_forward(self, x):
11 | return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps).to(x)
12 |
13 |
14 | LayerNorm.forward = LayerNorm_forward
15 | torch.nn.LayerNorm.forward = LayerNorm_forward
16 |
17 |
18 | def FP32LayerNorm_forward(self, x):
19 | origin_dtype = x.dtype
20 | return torch.nn.functional.layer_norm(
21 | x.float(),
22 | self.normalized_shape,
23 | self.weight.float() if self.weight is not None else None,
24 | self.bias.float() if self.bias is not None else None,
25 | self.eps,
26 | ).to(origin_dtype)
27 |
28 |
29 | FP32LayerNorm.forward = FP32LayerNorm_forward
30 |
31 |
32 | def RMSNorm_forward(self, hidden_states):
33 | input_dtype = hidden_states.dtype
34 | variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
35 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
36 |
37 | if self.weight is None:
38 | return hidden_states.to(input_dtype)
39 |
40 | return hidden_states.to(input_dtype) * self.weight.to(input_dtype)
41 |
42 |
43 | RMSNorm.forward = RMSNorm_forward
44 |
45 |
46 | def AdaLayerNormContinuous_forward(self, x, conditioning_embedding):
47 | emb = self.linear(self.silu(conditioning_embedding))
48 | scale, shift = emb.chunk(2, dim=1)
49 | x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
50 | return x
51 |
52 |
53 | AdaLayerNormContinuous.forward = AdaLayerNormContinuous_forward
54 |
--------------------------------------------------------------------------------
/diffusers_helper/k_diffusion/wrapper.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 |
4 | def append_dims(x, target_dims):
5 | return x[(...,) + (None,) * (target_dims - x.ndim)]
6 |
7 |
8 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
9 | if guidance_rescale == 0:
10 | return noise_cfg
11 |
12 | std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
13 | std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
14 | noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
15 | noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg
16 | return noise_cfg
17 |
18 |
19 | def fm_wrapper(transformer, t_scale=1000.0):
20 | def k_model(x, sigma, **extra_args):
21 | dtype = extra_args['dtype']
22 | cfg_scale = extra_args['cfg_scale']
23 | cfg_rescale = extra_args['cfg_rescale']
24 | concat_latent = extra_args['concat_latent']
25 |
26 | original_dtype = x.dtype
27 | sigma = sigma.float()
28 |
29 | x = x.to(dtype)
30 | timestep = (sigma * t_scale).to(dtype)
31 |
32 | if concat_latent is None:
33 | hidden_states = x
34 | else:
35 | hidden_states = torch.cat([x, concat_latent.to(x)], dim=1)
36 |
37 | pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float()
38 |
39 | if cfg_scale == 1.0:
40 | pred_negative = torch.zeros_like(pred_positive)
41 | else:
42 | pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float()
43 |
44 | pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative)
45 | pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale)
46 |
47 | x0 = x.float() - pred.float() * append_dims(sigma, x.ndim)
48 |
49 | return x0.to(dtype=original_dtype)
50 |
51 | return k_model
52 |
--------------------------------------------------------------------------------
/diffusers_helper/thread_utils.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | from threading import Thread, Lock
4 |
5 |
6 | class Listener:
7 | task_queue = []
8 | lock = Lock()
9 | thread = None
10 |
11 | @classmethod
12 | def _process_tasks(cls):
13 | while True:
14 | task = None
15 | with cls.lock:
16 | if cls.task_queue:
17 | task = cls.task_queue.pop(0)
18 |
19 | if task is None:
20 | time.sleep(0.001)
21 | continue
22 |
23 | func, args, kwargs = task
24 | try:
25 | func(*args, **kwargs)
26 | except Exception as e:
27 | print(f"Error in listener thread: {e}")
28 |
29 | @classmethod
30 | def add_task(cls, func, *args, **kwargs):
31 | with cls.lock:
32 | cls.task_queue.append((func, args, kwargs))
33 |
34 | if cls.thread is None:
35 | cls.thread = Thread(target=cls._process_tasks, daemon=True)
36 | cls.thread.start()
37 |
38 |
39 | def async_run(func, *args, **kwargs):
40 | Listener.add_task(func, *args, **kwargs)
41 |
42 |
43 | class FIFOQueue:
44 | def __init__(self):
45 | self.queue = []
46 | self.lock = Lock()
47 |
48 | def push(self, item):
49 | with self.lock:
50 | self.queue.append(item)
51 |
52 | def pop(self):
53 | with self.lock:
54 | if self.queue:
55 | return self.queue.pop(0)
56 | return None
57 |
58 | def top(self):
59 | with self.lock:
60 | if self.queue:
61 | return self.queue[0]
62 | return None
63 |
64 | def next(self):
65 | while True:
66 | with self.lock:
67 | if self.queue:
68 | return self.queue.pop(0)
69 |
70 | time.sleep(0.001)
71 |
72 |
73 | class AsyncStream:
74 | def __init__(self):
75 | self.input_queue = FIFOQueue()
76 | self.output_queue = FIFOQueue()
77 |
--------------------------------------------------------------------------------
/diffusers_helper/gradio/progress_bar.py:
--------------------------------------------------------------------------------
1 | progress_html = '''
2 |
9 | '''
10 |
11 | css = '''
12 | .loader-container {
13 | display: flex; /* Use flex to align items horizontally */
14 | align-items: center; /* Center items vertically within the container */
15 | white-space: nowrap; /* Prevent line breaks within the container */
16 | }
17 |
18 | .loader {
19 | border: 8px solid #f3f3f3; /* Light grey */
20 | border-top: 8px solid #3498db; /* Blue */
21 | border-radius: 50%;
22 | width: 30px;
23 | height: 30px;
24 | animation: spin 2s linear infinite;
25 | }
26 |
27 | @keyframes spin {
28 | 0% { transform: rotate(0deg); }
29 | 100% { transform: rotate(360deg); }
30 | }
31 |
32 | /* Style the progress bar */
33 | progress {
34 | appearance: none; /* Remove default styling */
35 | height: 20px; /* Set the height of the progress bar */
36 | border-radius: 5px; /* Round the corners of the progress bar */
37 | background-color: #f3f3f3; /* Light grey background */
38 | width: 100%;
39 | vertical-align: middle !important;
40 | }
41 |
42 | /* Style the progress bar container */
43 | .progress-container {
44 | margin-left: 20px;
45 | margin-right: 20px;
46 | flex-grow: 1; /* Allow the progress container to take up remaining space */
47 | }
48 |
49 | /* Set the color of the progress bar fill */
50 | progress::-webkit-progress-value {
51 | background-color: #3498db; /* Blue color for the fill */
52 | }
53 |
54 | progress::-moz-progress-bar {
55 | background-color: #3498db; /* Blue color for the fill in Firefox */
56 | }
57 |
58 | /* Style the text on the progress bar */
59 | progress::after {
60 | content: attr(value '%'); /* Display the progress value followed by '%' */
61 | position: absolute;
62 | top: 50%;
63 | left: 50%;
64 | transform: translate(-50%, -50%);
65 | color: white; /* Set text color */
66 | font-size: 14px; /* Set font size */
67 | }
68 |
69 | /* Style other texts */
70 | .loader-container > span {
71 | margin-left: 5px; /* Add spacing between the progress bar and the text */
72 | }
73 |
74 | .no-generating-animation > .generating {
75 | display: none !important;
76 | }
77 |
78 | '''
79 |
80 |
81 | def make_progress_bar_html(number, text):
82 | return progress_html.replace('*number*', str(number)).replace('*text*', text)
83 |
84 |
85 | def make_progress_bar_css():
86 | return css
87 |
--------------------------------------------------------------------------------
/diffusers_helper/hunyuan.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
4 | from diffusers_helper.utils import crop_or_pad_yield_mask
5 |
6 |
7 | @torch.no_grad()
8 | def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, max_length=256):
9 | assert isinstance(prompt, str)
10 |
11 | prompt = [prompt]
12 |
13 | # LLAMA
14 |
15 | prompt_llama = [DEFAULT_PROMPT_TEMPLATE["template"].format(p) for p in prompt]
16 | crop_start = DEFAULT_PROMPT_TEMPLATE["crop_start"]
17 |
18 | llama_inputs = tokenizer(
19 | prompt_llama,
20 | padding="max_length",
21 | max_length=max_length + crop_start,
22 | truncation=True,
23 | return_tensors="pt",
24 | return_length=False,
25 | return_overflowing_tokens=False,
26 | return_attention_mask=True,
27 | )
28 |
29 | llama_input_ids = llama_inputs.input_ids.to(text_encoder.device)
30 | llama_attention_mask = llama_inputs.attention_mask.to(text_encoder.device)
31 | llama_attention_length = int(llama_attention_mask.sum())
32 |
33 | llama_outputs = text_encoder(
34 | input_ids=llama_input_ids,
35 | attention_mask=llama_attention_mask,
36 | output_hidden_states=True,
37 | )
38 |
39 | llama_vec = llama_outputs.hidden_states[-3][:, crop_start:llama_attention_length]
40 | # llama_vec_remaining = llama_outputs.hidden_states[-3][:, llama_attention_length:]
41 | llama_attention_mask = llama_attention_mask[:, crop_start:llama_attention_length]
42 |
43 | assert torch.all(llama_attention_mask.bool())
44 |
45 | # CLIP
46 |
47 | clip_l_input_ids = tokenizer_2(
48 | prompt,
49 | padding="max_length",
50 | max_length=77,
51 | truncation=True,
52 | return_overflowing_tokens=False,
53 | return_length=False,
54 | return_tensors="pt",
55 | ).input_ids
56 | clip_l_pooler = text_encoder_2(clip_l_input_ids.to(text_encoder_2.device), output_hidden_states=False).pooler_output
57 |
58 | return llama_vec, clip_l_pooler
59 |
60 |
61 | @torch.no_grad()
62 | def vae_decode_fake(latents):
63 | latent_rgb_factors = [
64 | [-0.0395, -0.0331, 0.0445],
65 | [0.0696, 0.0795, 0.0518],
66 | [0.0135, -0.0945, -0.0282],
67 | [0.0108, -0.0250, -0.0765],
68 | [-0.0209, 0.0032, 0.0224],
69 | [-0.0804, -0.0254, -0.0639],
70 | [-0.0991, 0.0271, -0.0669],
71 | [-0.0646, -0.0422, -0.0400],
72 | [-0.0696, -0.0595, -0.0894],
73 | [-0.0799, -0.0208, -0.0375],
74 | [0.1166, 0.1627, 0.0962],
75 | [0.1165, 0.0432, 0.0407],
76 | [-0.2315, -0.1920, -0.1355],
77 | [-0.0270, 0.0401, -0.0821],
78 | [-0.0616, -0.0997, -0.0727],
79 | [0.0249, -0.0469, -0.1703]
80 | ] # From comfyui
81 |
82 | latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761]
83 |
84 | weight = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1)[:, :, None, None, None]
85 | bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype)
86 |
87 | images = torch.nn.functional.conv3d(latents, weight, bias=bias, stride=1, padding=0, dilation=1, groups=1)
88 | images = images.clamp(0.0, 1.0)
89 |
90 | return images
91 |
92 |
93 | @torch.no_grad()
94 | def vae_decode(latents, vae, image_mode=False):
95 | latents = latents / vae.config.scaling_factor
96 |
97 | if not image_mode:
98 | image = vae.decode(latents.to(device=vae.device, dtype=vae.dtype)).sample
99 | else:
100 | latents = latents.to(device=vae.device, dtype=vae.dtype).unbind(2)
101 | image = [vae.decode(l.unsqueeze(2)).sample for l in latents]
102 | image = torch.cat(image, dim=2)
103 |
104 | return image
105 |
106 |
107 | @torch.no_grad()
108 | def vae_encode(image, vae):
109 | latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()
110 | latents = latents * vae.config.scaling_factor
111 | return latents
112 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | hf_download/
2 | outputs/
3 | repo/
4 |
5 | # Byte-compiled / optimized / DLL files
6 | __pycache__/
7 | *.py[cod]
8 | *$py.class
9 |
10 | # C extensions
11 | *.so
12 |
13 | # Distribution / packaging
14 | .Python
15 | build/
16 | develop-eggs/
17 | dist/
18 | downloads/
19 | eggs/
20 | .eggs/
21 | lib/
22 | lib64/
23 | parts/
24 | sdist/
25 | var/
26 | wheels/
27 | share/python-wheels/
28 | *.egg-info/
29 | .installed.cfg
30 | *.egg
31 | MANIFEST
32 |
33 | # PyInstaller
34 | # Usually these files are written by a python script from a template
35 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
36 | *.manifest
37 | *.spec
38 |
39 | # Installer logs
40 | pip-log.txt
41 | pip-delete-this-directory.txt
42 |
43 | # Unit test / coverage reports
44 | htmlcov/
45 | .tox/
46 | .nox/
47 | .coverage
48 | .coverage.*
49 | .cache
50 | nosetests.xml
51 | coverage.xml
52 | *.cover
53 | *.py,cover
54 | .hypothesis/
55 | .pytest_cache/
56 | cover/
57 |
58 | # Translations
59 | *.mo
60 | *.pot
61 |
62 | # Django stuff:
63 | *.log
64 | local_settings.py
65 | db.sqlite3
66 | db.sqlite3-journal
67 |
68 | # Flask stuff:
69 | instance/
70 | .webassets-cache
71 |
72 | # Scrapy stuff:
73 | .scrapy
74 |
75 | # Sphinx documentation
76 | docs/_build/
77 |
78 | # PyBuilder
79 | .pybuilder/
80 | target/
81 |
82 | # Jupyter Notebook
83 | .ipynb_checkpoints
84 |
85 | # IPython
86 | profile_default/
87 | ipython_config.py
88 |
89 | # pyenv
90 | # For a library or package, you might want to ignore these files since the code is
91 | # intended to run in multiple environments; otherwise, check them in:
92 | # .python-version
93 |
94 | # pipenv
95 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
96 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
97 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
98 | # install all needed dependencies.
99 | #Pipfile.lock
100 |
101 | # UV
102 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
103 | # This is especially recommended for binary packages to ensure reproducibility, and is more
104 | # commonly ignored for libraries.
105 | #uv.lock
106 |
107 | # poetry
108 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
109 | # This is especially recommended for binary packages to ensure reproducibility, and is more
110 | # commonly ignored for libraries.
111 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
112 | #poetry.lock
113 |
114 | # pdm
115 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
116 | #pdm.lock
117 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
118 | # in version control.
119 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
120 | .pdm.toml
121 | .pdm-python
122 | .pdm-build/
123 |
124 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
125 | __pypackages__/
126 |
127 | # Celery stuff
128 | celerybeat-schedule
129 | celerybeat.pid
130 |
131 | # SageMath parsed files
132 | *.sage.py
133 |
134 | # Environments
135 | .env
136 | .venv
137 | env/
138 | venv/
139 | ENV/
140 | env.bak/
141 | venv.bak/
142 |
143 | # Spyder project settings
144 | .spyderproject
145 | .spyproject
146 |
147 | # Rope project settings
148 | .ropeproject
149 |
150 | # mkdocs documentation
151 | /site
152 |
153 | # mypy
154 | .mypy_cache/
155 | .dmypy.json
156 | dmypy.json
157 |
158 | # Pyre type checker
159 | .pyre/
160 |
161 | # pytype static type analyzer
162 | .pytype/
163 |
164 | # Cython debug symbols
165 | cython_debug/
166 |
167 | # PyCharm
168 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
169 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
170 | # and can be added to the global gitignore or merged into this file. For a more nuclear
171 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
172 | .idea/
173 |
174 | # Ruff stuff:
175 | .ruff_cache/
176 |
177 | # PyPI configuration file
178 | .pypirc
179 |
--------------------------------------------------------------------------------
/diffusers_helper/pipelines/k_diffusion_hunyuan.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import math
3 |
4 | from diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc
5 | from diffusers_helper.k_diffusion.wrapper import fm_wrapper
6 | from diffusers_helper.utils import repeat_to_batch_size
7 |
8 |
9 | def flux_time_shift(t, mu=1.15, sigma=1.0):
10 | return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
11 |
12 |
13 | def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, exp_max=7.0):
14 | k = (y2 - y1) / (x2 - x1)
15 | b = y1 - k * x1
16 | mu = k * context_length + b
17 | mu = min(mu, math.log(exp_max))
18 | return mu
19 |
20 |
21 | def get_flux_sigmas_from_mu(n, mu):
22 | sigmas = torch.linspace(1, 0, steps=n + 1)
23 | sigmas = flux_time_shift(sigmas, mu=mu)
24 | return sigmas
25 |
26 |
27 | @torch.inference_mode()
28 | def sample_hunyuan(
29 | transformer,
30 | sampler='unipc',
31 | initial_latent=None,
32 | concat_latent=None,
33 | strength=1.0,
34 | width=512,
35 | height=512,
36 | frames=16,
37 | real_guidance_scale=1.0,
38 | distilled_guidance_scale=6.0,
39 | guidance_rescale=0.0,
40 | shift=None,
41 | num_inference_steps=25,
42 | batch_size=None,
43 | generator=None,
44 | prompt_embeds=None,
45 | prompt_embeds_mask=None,
46 | prompt_poolers=None,
47 | negative_prompt_embeds=None,
48 | negative_prompt_embeds_mask=None,
49 | negative_prompt_poolers=None,
50 | dtype=torch.bfloat16,
51 | device=None,
52 | negative_kwargs=None,
53 | callback=None,
54 | **kwargs,
55 | ):
56 | device = device or transformer.device
57 |
58 | if batch_size is None:
59 | batch_size = int(prompt_embeds.shape[0])
60 |
61 | latents = torch.randn((batch_size, 16, (frames + 3) // 4, height // 8, width // 8), generator=generator, device=generator.device).to(device=device, dtype=torch.float32)
62 |
63 | B, C, T, H, W = latents.shape
64 | seq_length = T * H * W // 4
65 |
66 | if shift is None:
67 | mu = calculate_flux_mu(seq_length, exp_max=7.0)
68 | else:
69 | mu = math.log(shift)
70 |
71 | sigmas = get_flux_sigmas_from_mu(num_inference_steps, mu).to(device)
72 |
73 | k_model = fm_wrapper(transformer)
74 |
75 | if initial_latent is not None:
76 | sigmas = sigmas * strength
77 | first_sigma = sigmas[0].to(device=device, dtype=torch.float32)
78 | initial_latent = initial_latent.to(device=device, dtype=torch.float32)
79 | latents = initial_latent.float() * (1.0 - first_sigma) + latents.float() * first_sigma
80 |
81 | if concat_latent is not None:
82 | concat_latent = concat_latent.to(latents)
83 |
84 | distilled_guidance = torch.tensor([distilled_guidance_scale * 1000.0] * batch_size).to(device=device, dtype=dtype)
85 |
86 | prompt_embeds = repeat_to_batch_size(prompt_embeds, batch_size)
87 | prompt_embeds_mask = repeat_to_batch_size(prompt_embeds_mask, batch_size)
88 | prompt_poolers = repeat_to_batch_size(prompt_poolers, batch_size)
89 | negative_prompt_embeds = repeat_to_batch_size(negative_prompt_embeds, batch_size)
90 | negative_prompt_embeds_mask = repeat_to_batch_size(negative_prompt_embeds_mask, batch_size)
91 | negative_prompt_poolers = repeat_to_batch_size(negative_prompt_poolers, batch_size)
92 | concat_latent = repeat_to_batch_size(concat_latent, batch_size)
93 |
94 | sampler_kwargs = dict(
95 | dtype=dtype,
96 | cfg_scale=real_guidance_scale,
97 | cfg_rescale=guidance_rescale,
98 | concat_latent=concat_latent,
99 | positive=dict(
100 | pooled_projections=prompt_poolers,
101 | encoder_hidden_states=prompt_embeds,
102 | encoder_attention_mask=prompt_embeds_mask,
103 | guidance=distilled_guidance,
104 | **kwargs,
105 | ),
106 | negative=dict(
107 | pooled_projections=negative_prompt_poolers,
108 | encoder_hidden_states=negative_prompt_embeds,
109 | encoder_attention_mask=negative_prompt_embeds_mask,
110 | guidance=distilled_guidance,
111 | **(kwargs if negative_kwargs is None else {**kwargs, **negative_kwargs}),
112 | )
113 | )
114 |
115 | if sampler == 'unipc':
116 | results = sample_unipc(k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False, callback=callback)
117 | else:
118 | raise NotImplementedError(f'Sampler {sampler} is not supported.')
119 |
120 | return results
121 |
--------------------------------------------------------------------------------
/diffusers_helper/k_diffusion/uni_pc_fm.py:
--------------------------------------------------------------------------------
1 | # Better Flow Matching UniPC by Lvmin Zhang
2 | # (c) 2025
3 | # CC BY-SA 4.0
4 | # Attribution-ShareAlike 4.0 International Licence
5 |
6 |
7 | import torch
8 |
9 | from tqdm.auto import trange
10 |
11 |
12 | def expand_dims(v, dims):
13 | return v[(...,) + (None,) * (dims - 1)]
14 |
15 |
16 | class FlowMatchUniPC:
17 | def __init__(self, model, extra_args, variant='bh1'):
18 | self.model = model
19 | self.variant = variant
20 | self.extra_args = extra_args
21 |
22 | def model_fn(self, x, t):
23 | return self.model(x, t, **self.extra_args)
24 |
25 | def update_fn(self, x, model_prev_list, t_prev_list, t, order):
26 | assert order <= len(model_prev_list)
27 | dims = x.dim()
28 |
29 | t_prev_0 = t_prev_list[-1]
30 | lambda_prev_0 = - torch.log(t_prev_0)
31 | lambda_t = - torch.log(t)
32 | model_prev_0 = model_prev_list[-1]
33 |
34 | h = lambda_t - lambda_prev_0
35 |
36 | rks = []
37 | D1s = []
38 | for i in range(1, order):
39 | t_prev_i = t_prev_list[-(i + 1)]
40 | model_prev_i = model_prev_list[-(i + 1)]
41 | lambda_prev_i = - torch.log(t_prev_i)
42 | rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
43 | rks.append(rk)
44 | D1s.append((model_prev_i - model_prev_0) / rk)
45 |
46 | rks.append(1.)
47 | rks = torch.tensor(rks, device=x.device)
48 |
49 | R = []
50 | b = []
51 |
52 | hh = -h[0]
53 | h_phi_1 = torch.expm1(hh)
54 | h_phi_k = h_phi_1 / hh - 1
55 |
56 | factorial_i = 1
57 |
58 | if self.variant == 'bh1':
59 | B_h = hh
60 | elif self.variant == 'bh2':
61 | B_h = torch.expm1(hh)
62 | else:
63 | raise NotImplementedError('Bad variant!')
64 |
65 | for i in range(1, order + 1):
66 | R.append(torch.pow(rks, i - 1))
67 | b.append(h_phi_k * factorial_i / B_h)
68 | factorial_i *= (i + 1)
69 | h_phi_k = h_phi_k / hh - 1 / factorial_i
70 |
71 | R = torch.stack(R)
72 | b = torch.tensor(b, device=x.device)
73 |
74 | use_predictor = len(D1s) > 0
75 |
76 | if use_predictor:
77 | D1s = torch.stack(D1s, dim=1)
78 | if order == 2:
79 | rhos_p = torch.tensor([0.5], device=b.device)
80 | else:
81 | rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
82 | else:
83 | D1s = None
84 | rhos_p = None
85 |
86 | if order == 1:
87 | rhos_c = torch.tensor([0.5], device=b.device)
88 | else:
89 | rhos_c = torch.linalg.solve(R, b)
90 |
91 | x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
92 |
93 | if use_predictor:
94 | pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
95 | else:
96 | pred_res = 0
97 |
98 | x_t = x_t_ - expand_dims(B_h, dims) * pred_res
99 | model_t = self.model_fn(x_t, t)
100 |
101 | if D1s is not None:
102 | corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
103 | else:
104 | corr_res = 0
105 |
106 | D1_t = (model_t - model_prev_0)
107 | x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
108 |
109 | return x_t, model_t
110 |
111 | def sample(self, x, sigmas, callback=None, disable_pbar=False):
112 | order = min(3, len(sigmas) - 2)
113 | model_prev_list, t_prev_list = [], []
114 | for i in trange(len(sigmas) - 1, disable=disable_pbar):
115 | vec_t = sigmas[i].expand(x.shape[0])
116 |
117 | if i == 0:
118 | model_prev_list = [self.model_fn(x, vec_t)]
119 | t_prev_list = [vec_t]
120 | elif i < order:
121 | init_order = i
122 | x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
123 | model_prev_list.append(model_x)
124 | t_prev_list.append(vec_t)
125 | else:
126 | x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
127 | model_prev_list.append(model_x)
128 | t_prev_list.append(vec_t)
129 |
130 | model_prev_list = model_prev_list[-order:]
131 | t_prev_list = t_prev_list[-order:]
132 |
133 | if callback is not None:
134 | callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
135 |
136 | return model_prev_list[-1]
137 |
138 |
139 | def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
140 | assert variant in ['bh1', 'bh2']
141 | return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)
142 |
--------------------------------------------------------------------------------
/diffusers_helper/memory.py:
--------------------------------------------------------------------------------
1 | # By lllyasviel
2 |
3 |
4 | import torch
5 |
6 |
7 | # Detect available devices
8 | cpu = torch.device('cpu')
9 | if torch.cuda.is_available():
10 | gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
11 | elif torch.backends.mps.is_available():
12 | gpu = torch.device('mps')
13 | else:
14 | raise RuntimeError("No GPU device available. Please use a system with CUDA or MPS support.")
15 | gpu_complete_modules = []
16 |
17 |
18 | class DynamicSwapInstaller:
19 | @staticmethod
20 | def _install_module(module: torch.nn.Module, **kwargs):
21 | original_class = module.__class__
22 | module.__dict__['forge_backup_original_class'] = original_class
23 |
24 | def hacked_get_attr(self, name: str):
25 | if '_parameters' in self.__dict__:
26 | _parameters = self.__dict__['_parameters']
27 | if name in _parameters:
28 | p = _parameters[name]
29 | if p is None:
30 | return None
31 | if p.__class__ == torch.nn.Parameter:
32 | return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
33 | else:
34 | return p.to(**kwargs)
35 | if '_buffers' in self.__dict__:
36 | _buffers = self.__dict__['_buffers']
37 | if name in _buffers:
38 | return _buffers[name].to(**kwargs)
39 | return super(original_class, self).__getattr__(name)
40 |
41 | module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
42 | '__getattr__': hacked_get_attr,
43 | })
44 |
45 | return
46 |
47 | @staticmethod
48 | def _uninstall_module(module: torch.nn.Module):
49 | if 'forge_backup_original_class' in module.__dict__:
50 | module.__class__ = module.__dict__.pop('forge_backup_original_class')
51 | return
52 |
53 | @staticmethod
54 | def install_model(model: torch.nn.Module, **kwargs):
55 | for m in model.modules():
56 | DynamicSwapInstaller._install_module(m, **kwargs)
57 | return
58 |
59 | @staticmethod
60 | def uninstall_model(model: torch.nn.Module):
61 | for m in model.modules():
62 | DynamicSwapInstaller._uninstall_module(m)
63 | return
64 |
65 |
66 | def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
67 | if hasattr(model, 'scale_shift_table'):
68 | model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
69 | return
70 |
71 | for k, p in model.named_modules():
72 | if hasattr(p, 'weight'):
73 | p.to(target_device)
74 | return
75 |
76 |
77 | def get_cuda_free_memory_gb(device=None):
78 | if device is None:
79 | device = gpu
80 |
81 | if device.type == 'cuda':
82 | memory_stats = torch.cuda.memory_stats(device)
83 | bytes_active = memory_stats['active_bytes.all.current']
84 | bytes_reserved = memory_stats['reserved_bytes.all.current']
85 | bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
86 | bytes_inactive_reserved = bytes_reserved - bytes_active
87 | bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
88 | elif device.type == 'mps':
89 | # MPS doesn't provide detailed memory stats, return a best guess
90 | bytes_total_available = torch.mps.recommended_max_memory() - torch.mps.driver_allocated_memory()
91 |
92 | return bytes_total_available / (1024 ** 3)
93 |
94 |
95 | def empty_cache():
96 | if gpu.type == 'cuda':
97 | torch.cuda.empty_cache()
98 | elif gpu.type == 'mps':
99 | torch.mps.empty_cache()
100 |
101 | def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
102 | print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
103 |
104 | for m in model.modules():
105 | if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
106 | empty_cache()
107 | return
108 |
109 | if hasattr(m, 'weight'):
110 | m.to(device=target_device)
111 |
112 | model.to(device=target_device)
113 | empty_cache()
114 | return
115 |
116 |
117 | def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
118 | print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
119 |
120 | if target_device.type == 'cuda':
121 | for m in model.modules():
122 | if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
123 | empty_cache()
124 | return
125 |
126 | if hasattr(m, 'weight'):
127 | m.to(device=cpu)
128 | else:
129 | # For MPS, just move the model directly
130 | model.to(device=cpu)
131 |
132 | model.to(device=cpu)
133 | empty_cache()
134 | return
135 |
136 |
137 | def unload_complete_models(*args):
138 | for m in gpu_complete_modules + list(args):
139 | if m is not None:
140 | m.to(device=cpu)
141 | print(f'Unloaded {m.__class__.__name__} as complete.')
142 |
143 | gpu_complete_modules.clear()
144 | empty_cache()
145 | return
146 |
147 |
148 | def load_model_as_complete(model, target_device, unload=True):
149 | if unload:
150 | unload_complete_models()
151 |
152 | model.to(device=target_device)
153 | print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
154 |
155 | gpu_complete_modules.append(model)
156 | return
157 |
--------------------------------------------------------------------------------
/utils/lora_utils.py:
--------------------------------------------------------------------------------
1 | import os
2 | import torch
3 | from safetensors.torch import load_file
4 | from tqdm import tqdm
5 |
6 |
7 | def merge_lora_to_state_dict(
8 | state_dict: dict[str, torch.Tensor], lora_file: str, multiplier: float, device: torch.device
9 | ) -> dict[str, torch.Tensor]:
10 | """
11 | Merge LoRA weights into the state dict of a model.
12 | """
13 | lora_sd = load_file(lora_file)
14 |
15 | # Check the format of the LoRA file
16 | keys = list(lora_sd.keys())
17 | if keys[0].startswith("lora_unet_"):
18 | print(f"Musubi Tuner LoRA detected")
19 | return merge_musubi_tuner(lora_sd, state_dict, multiplier, device)
20 |
21 | transformer_prefixes = ["diffusion_model", "transformer"] # to ignore Text Encoder modules
22 | lora_suffix = None
23 | prefix = None
24 | for key in keys:
25 | if lora_suffix is None and "lora_A" in key:
26 | lora_suffix = "lora_A"
27 | if prefix is None:
28 | pfx = key.split(".")[0]
29 | if pfx in transformer_prefixes:
30 | prefix = pfx
31 | if lora_suffix is not None and prefix is not None:
32 | break
33 |
34 | if lora_suffix == "lora_A" and prefix is not None:
35 | print(f"Diffusion-pipe (?) LoRA detected")
36 | return merge_diffusion_pipe_or_something(lora_sd, state_dict, "lora_unet_", multiplier, device)
37 |
38 | print(f"LoRA file format not recognized: {os.path.basename(lora_file)}")
39 | return state_dict
40 |
41 |
42 | def merge_diffusion_pipe_or_something(
43 | lora_sd: dict[str, torch.Tensor], state_dict: dict[str, torch.Tensor], prefix: str, multiplier: float, device: torch.device
44 | ) -> dict[str, torch.Tensor]:
45 | """
46 | Convert LoRA weights to the format used by the diffusion pipeline to Musubi Tuner.
47 | Copy from Musubi Tuner repo.
48 | """
49 | # convert from diffusers(?) to default LoRA
50 | # Diffusers format: {"diffusion_model.module.name.lora_A.weight": weight, "diffusion_model.module.name.lora_B.weight": weight, ...}
51 | # default LoRA format: {"prefix_module_name.lora_down.weight": weight, "prefix_module_name.lora_up.weight": weight, ...}
52 |
53 | # note: Diffusers has no alpha, so alpha is set to rank
54 | new_weights_sd = {}
55 | lora_dims = {}
56 | for key, weight in lora_sd.items():
57 | diffusers_prefix, key_body = key.split(".", 1)
58 | if diffusers_prefix != "diffusion_model" and diffusers_prefix != "transformer":
59 | print(f"unexpected key: {key} in diffusers format")
60 | continue
61 |
62 | new_key = f"{prefix}{key_body}".replace(".", "_").replace("_lora_A_", ".lora_down.").replace("_lora_B_", ".lora_up.")
63 | new_weights_sd[new_key] = weight
64 |
65 | lora_name = new_key.split(".")[0] # before first dot
66 | if lora_name not in lora_dims and "lora_down" in new_key:
67 | lora_dims[lora_name] = weight.shape[0]
68 |
69 | # add alpha with rank
70 | for lora_name, dim in lora_dims.items():
71 | new_weights_sd[f"{lora_name}.alpha"] = torch.tensor(dim)
72 |
73 | return merge_musubi_tuner(new_weights_sd, state_dict, multiplier, device)
74 |
75 |
76 | def merge_musubi_tuner(
77 | lora_sd: dict[str, torch.Tensor], state_dict: dict[str, torch.Tensor], multiplier: float, device: torch.device
78 | ) -> dict[str, torch.Tensor]:
79 | """
80 | Merge LoRA weights into the state dict of a model.
81 | """
82 | # Check LoRA is for FramePack or for HunyuanVideo
83 | is_hunyuan = False
84 | for key in lora_sd.keys():
85 | if "double_blocks" in key or "single_blocks" in key:
86 | is_hunyuan = True
87 | break
88 | if is_hunyuan:
89 | print("HunyuanVideo LoRA detected, converting to FramePack format")
90 | lora_sd = convert_hunyuan_to_framepack(lora_sd)
91 |
92 | # Merge LoRA weights into the state dict
93 | print(f"Merging LoRA weights into state dict. multiplier: {multiplier}")
94 |
95 | # Create module map
96 | name_to_original_key = {}
97 | for key in state_dict.keys():
98 | if key.endswith(".weight"):
99 | lora_name = key.rsplit(".", 1)[0] # remove trailing ".weight"
100 | lora_name = "lora_unet_" + lora_name.replace(".", "_")
101 | if lora_name not in name_to_original_key:
102 | name_to_original_key[lora_name] = key
103 |
104 | # Merge LoRA weights
105 | keys = list([k for k in lora_sd.keys() if "lora_down" in k])
106 | for key in tqdm(keys, desc="Merging LoRA weights"):
107 | up_key = key.replace("lora_down", "lora_up")
108 | alpha_key = key[: key.index("lora_down")] + "alpha"
109 |
110 | # find original key for this lora
111 | module_name = ".".join(key.split(".")[:-2]) # remove trailing ".lora_down.weight"
112 | if module_name not in name_to_original_key:
113 | print(f"No module found for LoRA weight: {key}")
114 | continue
115 |
116 | original_key = name_to_original_key[module_name]
117 |
118 | down_weight = lora_sd[key]
119 | up_weight = lora_sd[up_key]
120 |
121 | dim = down_weight.size()[0]
122 | alpha = lora_sd.get(alpha_key, dim)
123 | scale = alpha / dim
124 |
125 | weight = state_dict[original_key]
126 | original_device = weight.device
127 | if original_device != device:
128 | weight = weight.to(device) # to make calculation faster
129 |
130 | down_weight = down_weight.to(device)
131 | up_weight = up_weight.to(device)
132 |
133 | # W <- W + U * D
134 | if len(weight.size()) == 2:
135 | # linear
136 | if len(up_weight.size()) == 4: # use linear projection mismatch
137 | up_weight = up_weight.squeeze(3).squeeze(2)
138 | down_weight = down_weight.squeeze(3).squeeze(2)
139 | weight = weight + multiplier * (up_weight @ down_weight) * scale
140 | elif down_weight.size()[2:4] == (1, 1):
141 | # conv2d 1x1
142 | weight = (
143 | weight
144 | + multiplier
145 | * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
146 | * scale
147 | )
148 | else:
149 | # conv2d 3x3
150 | conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
151 | # logger.info(conved.size(), weight.size(), module.stride, module.padding)
152 | weight = weight + multiplier * conved * scale
153 |
154 | weight = weight.to(original_device) # move back to original device
155 | state_dict[original_key] = weight
156 |
157 | return state_dict
158 |
159 |
160 | def convert_hunyuan_to_framepack(lora_sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
161 | """
162 | Convert HunyuanVideo LoRA weights to FramePack format.
163 | """
164 | new_lora_sd = {}
165 | for key, weight in lora_sd.items():
166 | if "double_blocks" in key:
167 | key = key.replace("double_blocks", "transformer_blocks")
168 | key = key.replace("img_mod_linear", "norm1_linear")
169 | key = key.replace("img_attn_qkv", "attn_to_QKV") # split later
170 | key = key.replace("img_attn_proj", "attn_to_out_0")
171 | key = key.replace("img_mlp_fc1", "ff_net_0_proj")
172 | key = key.replace("img_mlp_fc2", "ff_net_2")
173 | key = key.replace("txt_mod_linear", "norm1_context_linear")
174 | key = key.replace("txt_attn_qkv", "attn_add_QKV_proj") # split later
175 | key = key.replace("txt_attn_proj", "attn_to_add_out")
176 | key = key.replace("txt_mlp_fc1", "ff_context_net_0_proj")
177 | key = key.replace("txt_mlp_fc2", "ff_context_net_2")
178 | elif "single_blocks" in key:
179 | key = key.replace("single_blocks", "single_transformer_blocks")
180 | key = key.replace("linear1", "attn_to_QKVM") # split later
181 | key = key.replace("linear2", "proj_out")
182 | key = key.replace("modulation_linear", "norm_linear")
183 | else:
184 | print(f"Unsupported module name: {key}, only double_blocks and single_blocks are supported")
185 | continue
186 |
187 | if "QKVM" in key:
188 | # split QKVM into Q, K, V, M
189 | key_q = key.replace("QKVM", "q")
190 | key_k = key.replace("QKVM", "k")
191 | key_v = key.replace("QKVM", "v")
192 | key_m = key.replace("attn_to_QKVM", "proj_mlp")
193 | if "_down" in key or "alpha" in key:
194 | # copy QKVM weight or alpha to Q, K, V, M
195 | assert "alpha" in key or weight.size(1) == 3072, f"QKVM weight size mismatch: {key}. {weight.size()}"
196 | new_lora_sd[key_q] = weight
197 | new_lora_sd[key_k] = weight
198 | new_lora_sd[key_v] = weight
199 | new_lora_sd[key_m] = weight
200 | elif "_up" in key:
201 | # split QKVM weight into Q, K, V, M
202 | assert weight.size(0) == 21504, f"QKVM weight size mismatch: {key}. {weight.size()}"
203 | new_lora_sd[key_q] = weight[:3072]
204 | new_lora_sd[key_k] = weight[3072 : 3072 * 2]
205 | new_lora_sd[key_v] = weight[3072 * 2 : 3072 * 3]
206 | new_lora_sd[key_m] = weight[3072 * 3 :] # 21504 - 3072 * 3 = 12288
207 | else:
208 | print(f"Unsupported module name: {key}")
209 | continue
210 | elif "QKV" in key:
211 | # split QKV into Q, K, V
212 | key_q = key.replace("QKV", "q")
213 | key_k = key.replace("QKV", "k")
214 | key_v = key.replace("QKV", "v")
215 | if "_down" in key or "alpha" in key:
216 | # copy QKV weight or alpha to Q, K, V
217 | assert "alpha" in key or weight.size(1) == 3072, f"QKV weight size mismatch: {key}. {weight.size()}"
218 | new_lora_sd[key_q] = weight
219 | new_lora_sd[key_k] = weight
220 | new_lora_sd[key_v] = weight
221 | elif "_up" in key:
222 | # split QKV weight into Q, K, V
223 | assert weight.size(0) == 3072 * 3, f"QKV weight size mismatch: {key}. {weight.size()}"
224 | new_lora_sd[key_q] = weight[:3072]
225 | new_lora_sd[key_k] = weight[3072 : 3072 * 2]
226 | new_lora_sd[key_v] = weight[3072 * 2 :]
227 | else:
228 | print(f"Unsupported module name: {key}")
229 | continue
230 | else:
231 | # no split needed
232 | new_lora_sd[key] = weight
233 |
234 | return new_lora_sd
235 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Apache License
2 | Version 2.0, January 2004
3 | http://www.apache.org/licenses/
4 |
5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6 |
7 | 1. Definitions.
8 |
9 | "License" shall mean the terms and conditions for use, reproduction,
10 | and distribution as defined by Sections 1 through 9 of this document.
11 |
12 | "Licensor" shall mean the copyright owner or entity authorized by
13 | the copyright owner that is granting the License.
14 |
15 | "Legal Entity" shall mean the union of the acting entity and all
16 | other entities that control, are controlled by, or are under common
17 | control with that entity. For the purposes of this definition,
18 | "control" means (i) the power, direct or indirect, to cause the
19 | direction or management of such entity, whether by contract or
20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
21 | outstanding shares, or (iii) beneficial ownership of such entity.
22 |
23 | "You" (or "Your") shall mean an individual or Legal Entity
24 | exercising permissions granted by this License.
25 |
26 | "Source" form shall mean the preferred form for making modifications,
27 | including but not limited to software source code, documentation
28 | source, and configuration files.
29 |
30 | "Object" form shall mean any form resulting from mechanical
31 | transformation or translation of a Source form, including but
32 | not limited to compiled object code, generated documentation,
33 | and conversions to other media types.
34 |
35 | "Work" shall mean the work of authorship, whether in Source or
36 | Object form, made available under the License, as indicated by a
37 | copyright notice that is included in or attached to the work
38 | (an example is provided in the Appendix below).
39 |
40 | "Derivative Works" shall mean any work, whether in Source or Object
41 | form, that is based on (or derived from) the Work and for which the
42 | editorial revisions, annotations, elaborations, or other modifications
43 | represent, as a whole, an original work of authorship. For the purposes
44 | of this License, Derivative Works shall not include works that remain
45 | separable from, or merely link (or bind by name) to the interfaces of,
46 | the Work and Derivative Works thereof.
47 |
48 | "Contribution" shall mean any work of authorship, including
49 | the original version of the Work and any modifications or additions
50 | to that Work or Derivative Works thereof, that is intentionally
51 | submitted to Licensor for inclusion in the Work by the copyright owner
52 | or by an individual or Legal Entity authorized to submit on behalf of
53 | the copyright owner. For the purposes of this definition, "submitted"
54 | means any form of electronic, verbal, or written communication sent
55 | to the Licensor or its representatives, including but not limited to
56 | communication on electronic mailing lists, source code control systems,
57 | and issue tracking systems that are managed by, or on behalf of, the
58 | Licensor for the purpose of discussing and improving the Work, but
59 | excluding communication that is conspicuously marked or otherwise
60 | designated in writing by the copyright owner as "Not a Contribution."
61 |
62 | "Contributor" shall mean Licensor and any individual or Legal Entity
63 | on behalf of whom a Contribution has been received by Licensor and
64 | subsequently incorporated within the Work.
65 |
66 | 2. Grant of Copyright License. Subject to the terms and conditions of
67 | this License, each Contributor hereby grants to You a perpetual,
68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69 | copyright license to reproduce, prepare Derivative Works of,
70 | publicly display, publicly perform, sublicense, and distribute the
71 | Work and such Derivative Works in Source or Object form.
72 |
73 | 3. Grant of Patent License. Subject to the terms and conditions of
74 | this License, each Contributor hereby grants to You a perpetual,
75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76 | (except as stated in this section) patent license to make, have made,
77 | use, offer to sell, sell, import, and otherwise transfer the Work,
78 | where such license applies only to those patent claims licensable
79 | by such Contributor that are necessarily infringed by their
80 | Contribution(s) alone or by combination of their Contribution(s)
81 | with the Work to which such Contribution(s) was submitted. If You
82 | institute patent litigation against any entity (including a
83 | cross-claim or counterclaim in a lawsuit) alleging that the Work
84 | or a Contribution incorporated within the Work constitutes direct
85 | or contributory patent infringement, then any patent licenses
86 | granted to You under this License for that Work shall terminate
87 | as of the date such litigation is filed.
88 |
89 | 4. Redistribution. You may reproduce and distribute copies of the
90 | Work or Derivative Works thereof in any medium, with or without
91 | modifications, and in Source or Object form, provided that You
92 | meet the following conditions:
93 |
94 | (a) You must give any other recipients of the Work or
95 | Derivative Works a copy of this License; and
96 |
97 | (b) You must cause any modified files to carry prominent notices
98 | stating that You changed the files; and
99 |
100 | (c) You must retain, in the Source form of any Derivative Works
101 | that You distribute, all copyright, patent, trademark, and
102 | attribution notices from the Source form of the Work,
103 | excluding those notices that do not pertain to any part of
104 | the Derivative Works; and
105 |
106 | (d) If the Work includes a "NOTICE" text file as part of its
107 | distribution, then any Derivative Works that You distribute must
108 | include a readable copy of the attribution notices contained
109 | within such NOTICE file, excluding those notices that do not
110 | pertain to any part of the Derivative Works, in at least one
111 | of the following places: within a NOTICE text file distributed
112 | as part of the Derivative Works; within the Source form or
113 | documentation, if provided along with the Derivative Works; or,
114 | within a display generated by the Derivative Works, if and
115 | wherever such third-party notices normally appear. The contents
116 | of the NOTICE file are for informational purposes only and
117 | do not modify the License. You may add Your own attribution
118 | notices within Derivative Works that You distribute, alongside
119 | or as an addendum to the NOTICE text from the Work, provided
120 | that such additional attribution notices cannot be construed
121 | as modifying the License.
122 |
123 | You may add Your own copyright statement to Your modifications and
124 | may provide additional or different license terms and conditions
125 | for use, reproduction, or distribution of Your modifications, or
126 | for any such Derivative Works as a whole, provided Your use,
127 | reproduction, and distribution of the Work otherwise complies with
128 | the conditions stated in this License.
129 |
130 | 5. Submission of Contributions. Unless You explicitly state otherwise,
131 | any Contribution intentionally submitted for inclusion in the Work
132 | by You to the Licensor shall be under the terms and conditions of
133 | this License, without any additional terms or conditions.
134 | Notwithstanding the above, nothing herein shall supersede or modify
135 | the terms of any separate license agreement you may have executed
136 | with Licensor regarding such Contributions.
137 |
138 | 6. Trademarks. This License does not grant permission to use the trade
139 | names, trademarks, service marks, or product names of the Licensor,
140 | except as required for reasonable and customary use in describing the
141 | origin of the Work and reproducing the content of the NOTICE file.
142 |
143 | 7. Disclaimer of Warranty. Unless required by applicable law or
144 | agreed to in writing, Licensor provides the Work (and each
145 | Contributor provides its Contributions) on an "AS IS" BASIS,
146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 | implied, including, without limitation, any warranties or conditions
148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 | PARTICULAR PURPOSE. You are solely responsible for determining the
150 | appropriateness of using or redistributing the Work and assume any
151 | risks associated with Your exercise of permissions under this License.
152 |
153 | 8. Limitation of Liability. In no event and under no legal theory,
154 | whether in tort (including negligence), contract, or otherwise,
155 | unless required by applicable law (such as deliberate and grossly
156 | negligent acts) or agreed to in writing, shall any Contributor be
157 | liable to You for damages, including any direct, indirect, special,
158 | incidental, or consequential damages of any character arising as a
159 | result of this License or out of the use or inability to use the
160 | Work (including but not limited to damages for loss of goodwill,
161 | work stoppage, computer failure or malfunction, or any and all
162 | other commercial damages or losses), even if such Contributor
163 | has been advised of the possibility of such damages.
164 |
165 | 9. Accepting Warranty or Additional Liability. While redistributing
166 | the Work or Derivative Works thereof, You may choose to offer,
167 | and charge a fee for, acceptance of support, warranty, indemnity,
168 | or other liability obligations and/or rights consistent with this
169 | License. However, in accepting such obligations, You may act only
170 | on Your own behalf and on Your sole responsibility, not on behalf
171 | of any other Contributor, and only if You agree to indemnify,
172 | defend, and hold each Contributor harmless for any liability
173 | incurred by, or claims asserted against, such Contributor by reason
174 | of your accepting any such warranty or additional liability.
175 |
176 | END OF TERMS AND CONDITIONS
177 |
178 | APPENDIX: How to apply the Apache License to your work.
179 |
180 | To apply the Apache License to your work, attach the following
181 | boilerplate notice, with the fields enclosed by brackets "[]"
182 | replaced with your own identifying information. (Don't include
183 | the brackets!) The text should be enclosed in the appropriate
184 | comment syntax for the file format. We also recommend that a
185 | file or class name and description of purpose be included on the
186 | same "printed page" as the copyright notice for easier
187 | identification within third-party archives.
188 |
189 | Copyright [yyyy] [name of copyright owner]
190 |
191 | Licensed under the Apache License, Version 2.0 (the "License");
192 | you may not use this file except in compliance with the License.
193 | You may obtain a copy of the License at
194 |
195 | http://www.apache.org/licenses/LICENSE-2.0
196 |
197 | Unless required by applicable law or agreed to in writing, software
198 | distributed under the License is distributed on an "AS IS" BASIS,
199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 | See the License for the specific language governing permissions and
201 | limitations under the License.
202 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 | 
3 |
4 |
5 | # FramePack
6 |
7 | Official implementation and desktop software for ["Packing Input Frame Context in Next-Frame Prediction Models for Video Generation"](https://lllyasviel.github.io/frame_pack_gitpage/).
8 |
9 | Links: [**Paper**](https://arxiv.org/abs/2504.12626), [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/)
10 |
11 | FramePack is a next-frame (next-frame-section) prediction neural network structure that generates videos progressively.
12 |
13 | FramePack compresses input contexts to a constant length so that the generation workload is invariant to video length.
14 |
15 | FramePack can process a very large number of frames with 13B models even on laptop GPUs.
16 |
17 | FramePack can be trained with a much larger batch size, similar to the batch size for image diffusion training.
18 |
19 | This version of FramePack is rendering videos at 24fps.
20 |
21 | **Video diffusion, but feels like image diffusion.**
22 |
23 | # News
24 |
25 | **2025 May 03:** The FramePack-F1 is released. [Try it here.](https://github.com/lllyasviel/FramePack/discussions/459)
26 |
27 | Note that this GitHub repository is the only official FramePack website. We do not have any web services. All other websites are spam and fake, including but not limited to `framepack.co`, `frame_pack.co`, `framepack.net`, `frame_pack.net`, `framepack.ai`, `frame_pack.ai`, `framepack.pro`, `frame_pack.pro`, `framepack.cc`, `frame_pack.cc`,`framepackai.co`, `frame_pack_ai.co`, `framepackai.net`, `frame_pack_ai.net`, `framepackai.pro`, `frame_pack_ai.pro`, `framepackai.cc`, `frame_pack_ai.cc`, and so on. Again, they are all spam and fake. **Do not pay money or download files from any of those websites.**
28 |
29 | # Requirements
30 |
31 | Note that this repo is a functional desktop software with minimal standalone high-quality sampling system and memory management.
32 |
33 | **Start with this repo before you try anything else!**
34 |
35 | Requirements:
36 |
37 | * Nvidia GPU in RTX 30XX, 40XX, 50XX series that supports fp16 and bf16. The GTX 10XX/20XX are not tested.
38 | * Linux or Windows operating system.
39 | * At least 6GB GPU memory.
40 |
41 | To generate 1-minute video (60 seconds) at 30fps (1800 frames) using 13B model, the minimal required GPU memory is 6GB. (Yes 6 GB, not a typo. Laptop GPUs are okay.)
42 |
43 | About speed, on my RTX 4090 desktop it generates at a speed of 2.5 seconds/frame (unoptimized) or 1.5 seconds/frame (teacache). On my laptops like 3070ti laptop or 3060 laptop, it is about 4x to 8x slower. [Troubleshoot if your speed is much slower than this.](https://github.com/lllyasviel/FramePack/issues/151#issuecomment-2817054649)
44 |
45 | In any case, you will directly see the generated frames since it is next-frame(-section) prediction. So you will get lots of visual feedback before the entire video is generated.
46 |
47 | # Installation
48 |
49 | **Windows**:
50 |
51 | [>>> Click Here to Download One-Click Package (CUDA 12.6 + Pytorch 2.6) <<<](https://github.com/lllyasviel/FramePack/releases/download/windows/framepack_cu126_torch26.7z)
52 |
53 | After you download, you uncompress, use `update.bat` to update, and use `run.bat` to run.
54 |
55 | Note that running `update.bat` is important, otherwise you may be using a previous version with potential bugs unfixed.
56 |
57 | 
58 |
59 | Note that the models will be downloaded automatically. You will download more than 30GB from HuggingFace.
60 |
61 | **Linux**:
62 |
63 | We recommend having an independent Python 3.10.
64 |
65 | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
66 | pip install -r requirements.txt
67 |
68 | To start the GUI, run:
69 |
70 | python demo_gradio.py
71 |
72 | Note that it supports `--share`, `--port`, `--server`, and so on.
73 |
74 | The software supports PyTorch attention, xformers, flash-attn, sage-attention. By default, it will just use PyTorch attention. You can install those attention kernels if you know how.
75 |
76 | For example, to install sage-attention (linux):
77 |
78 | pip install sageattention==1.0.6
79 |
80 | However, you are highly recommended to first try without sage-attention since it will influence results, though the influence is minimal.
81 |
82 | To start the GUI, run:
83 |
84 | python demo_gradio.py
85 |
86 | Note that it supports `--share`, `--port`, `--server`, and so on.
87 |
88 | **macOS**:
89 |
90 | FramePack recommends using Python 3.10. If you have [homebrew](https://brew.sh/) installed, you can install Python 3.10 using brew.
91 | ```bash
92 | brew install python@3.10
93 | ```
94 |
95 | To install dependencies
96 | ```bash
97 | pip3.10 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
98 | pip3.10 install -r requirements.txt
99 | ```
100 |
101 | # Starting FramePack on macOS
102 | To start the GUI, run:
103 |
104 | python3.10 demo_gradio.py
105 |
106 | Some additional arguments you can specify are:
107 | * `--share`: Enable sharing the Gradio interface via a public URL
108 | * `--server`: Specify the server address (default: '0.0.0.0')
109 | * `--port`: Specify the port number to run the server on (default: 7860)
110 | * `--output_dir`: Set the directory for saving generated outputs (default: './outputs')
111 |
112 | This version now defaults to 416, but may also be successful with slightly higher resolutions like 480. Feel free to experiment with the setting.
113 | (NOTE: The 640 value means it will generate a video that is roughly 640x640 for a square image. If the image has a different aspect ratio, the resolution will attempt to match the ratio and keep roughly the same number of pixels.)
114 |
115 | # GUI
116 |
117 | 
118 |
119 | On the left you upload an image and write a prompt.
120 |
121 | On the right are the generated videos and latent previews.
122 |
123 | Because this is a next-frame-section prediction model, videos will be generated longer and longer.
124 |
125 | You will see the progress bar for each section and the latent preview for the next section.
126 |
127 | Note that the initial progress may be slower than later diffusion as the device may need some warmup.
128 |
129 | # Sanity Check
130 |
131 | Before trying your own inputs, we highly recommend going through the sanity check to find out if any hardware or software went wrong.
132 |
133 | Next-frame-section prediction models are very sensitive to subtle differences in noise and hardware. Usually, people will get slightly different results on different devices, but the results should look overall similar. In some cases, if possible, you'll get exactly the same results.
134 |
135 | ## Image-to-5-seconds
136 |
137 | Download this image:
138 |
139 | 
140 |
141 | Copy this prompt:
142 |
143 | `The man dances energetically, leaping mid-air with fluid arm swings and quick footwork.`
144 |
145 | Set like this:
146 |
147 | (all default parameters, with teacache turned off)
148 | 
149 |
150 | The result will be:
151 |
152 |
153 |
154 | |
155 |
160 | |
161 |
162 |
163 | |
164 | Video may be compressed by GitHub
165 | |
166 |
167 |
168 |
169 | **Important Note:**
170 |
171 | Again, this is a next-frame-section prediction model. This means you will generate videos frame-by-frame or section-by-section.
172 |
173 | **If you get a much shorter video in the UI, like a video with only 1 second, then it is totally expected.** You just need to wait. More sections will be generated to complete the video.
174 |
175 | ## Know the influence of TeaCache and Quantization
176 |
177 | Download this image:
178 |
179 | 
180 |
181 | Copy this prompt:
182 |
183 | `The girl dances gracefully, with clear movements, full of charm.`
184 |
185 | Set like this:
186 |
187 | 
188 |
189 | Turn off teacache:
190 |
191 | 
192 |
193 | You will get this:
194 |
195 |
196 |
197 | |
198 |
203 | |
204 |
205 |
206 | |
207 | Video may be compressed by GitHub
208 | |
209 |
210 |
211 |
212 | Now turn on teacache:
213 |
214 | 
215 |
216 | About 30% users will get this (the other 70% will get other random results depending on their hardware):
217 |
218 |
219 |
220 | |
221 |
226 | |
227 |
228 |
229 | |
230 | A typical worse result.
231 | |
232 |
233 |
234 |
235 | So you can see that teacache is not really lossless and sometimes can influence the result a lot.
236 |
237 | We recommend using teacache to try ideas and then using the full diffusion process to get high-quality results.
238 |
239 | This recommendation also applies to sage-attention, bnb quant, gguf, etc., etc.
240 |
241 | ## Image-to-1-minute
242 |
243 | 
244 |
245 | `The girl dances gracefully, with clear movements, full of charm.`
246 |
247 | 
248 |
249 | Set video length to 60 seconds:
250 |
251 | 
252 |
253 | If everything is in order you will get some result like this eventually.
254 |
255 | 60s version:
256 |
257 |
258 |
259 | |
260 |
265 | |
266 |
267 |
268 | |
269 | Video may be compressed by GitHub
270 | |
271 |
272 |
273 |
274 | 6s version:
275 |
276 |
277 |
278 | |
279 |
284 | |
285 |
286 |
287 | |
288 | Video may be compressed by GitHub
289 | |
290 |
291 |
292 |
293 | # More Examples
294 |
295 | Many more examples are in [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/).
296 |
297 | Below are some more examples that you may be interested in reproducing.
298 |
299 | ---
300 |
301 | 
302 |
303 | `The girl dances gracefully, with clear movements, full of charm.`
304 |
305 | 
306 |
307 |
308 |
309 | |
310 |
315 | |
316 |
317 |
318 | |
319 | Video may be compressed by GitHub
320 | |
321 |
322 |
323 |
324 | ---
325 |
326 | 
327 |
328 | `The girl suddenly took out a sign that said “cute” using right hand`
329 |
330 | 
331 |
332 |
333 |
334 | |
335 |
340 | |
341 |
342 |
343 | |
344 | Video may be compressed by GitHub
345 | |
346 |
347 |
348 |
349 | ---
350 |
351 | 
352 |
353 | `The girl skateboarding, repeating the endless spinning and dancing and jumping on a skateboard, with clear movements, full of charm.`
354 |
355 | 
356 |
357 |
358 |
359 | |
360 |
365 | |
366 |
367 |
368 | |
369 | Video may be compressed by GitHub
370 | |
371 |
372 |
373 |
374 | ---
375 |
376 | 
377 |
378 | `The girl dances gracefully, with clear movements, full of charm.`
379 |
380 | 
381 |
382 |
383 |
384 | |
385 |
390 | |
391 |
392 |
393 | |
394 | Video may be compressed by GitHub
395 | |
396 |
397 |
398 |
399 | ---
400 |
401 | 
402 |
403 | `The man dances flamboyantly, swinging his hips and striking bold poses with dramatic flair.`
404 |
405 | 
406 |
407 |
408 |
409 | |
410 |
415 | |
416 |
417 |
418 | |
419 | Video may be compressed by GitHub
420 | |
421 |
422 |
423 |
424 | ---
425 |
426 | 
427 |
428 | `The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements.`
429 |
430 | 
431 |
432 |
433 |
434 |
435 | |
436 |
441 | |
442 |
443 |
444 | |
445 | Video may be compressed by GitHub
446 | |
447 |
448 |
449 |
450 | ---
451 |
452 | 
453 |
454 | `The young man writes intensely, flipping papers and adjusting his glasses with swift, focused movements.`
455 |
456 | 
457 |
458 |
459 |
460 | |
461 |
466 | |
467 |
468 |
469 | |
470 | Video may be compressed by GitHub
471 | |
472 |
473 |
474 |
475 | ---
476 |
477 | # Prompting Guideline
478 |
479 | Many people would ask how to write better prompts.
480 |
481 | Below is a ChatGPT template that I personally often use to get prompts:
482 |
483 | You are an assistant that writes short, motion-focused prompts for animating images.
484 |
485 | When the user sends an image, respond with a single, concise prompt describing visual motion (such as human activity, moving objects, or camera movements). Focus only on how the scene could come alive and become dynamic using brief phrases.
486 |
487 | Larger and more dynamic motions (like dancing, jumping, running, etc.) are preferred over smaller or more subtle ones (like standing still, sitting, etc.).
488 |
489 | Describe subject, then motion, then other things. For example: "The girl dances gracefully, with clear movements, full of charm."
490 |
491 | If there is something that can dance (like a man, girl, robot, etc.), then prefer to describe it as dancing.
492 |
493 | Stay in a loop: one image in, one motion prompt out. Do not explain, ask questions, or generate multiple options.
494 |
495 | You paste the instruct to ChatGPT and then feed it an image to get prompt like this:
496 |
497 | 
498 |
499 | *The man dances powerfully, striking sharp poses and gliding smoothly across the reflective floor.*
500 |
501 | Usually this will give you a prompt that works well.
502 |
503 | You can also write prompts yourself. Concise prompts are usually preferred, for example:
504 |
505 | *The girl dances gracefully, with clear movements, full of charm.*
506 |
507 | *The man dances powerfully, with clear movements, full of energy.*
508 |
509 | and so on.
510 |
511 | # Cite
512 |
513 | @article{zhang2025framepack,
514 | title={Packing Input Frame Contexts in Next-Frame Prediction Models for Video Generation},
515 | author={Lvmin Zhang and Maneesh Agrawala},
516 | journal={Arxiv},
517 | year={2025}
518 | }
519 |
--------------------------------------------------------------------------------
/diffusers_helper/utils.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import json
4 | import random
5 | import glob
6 | import torch
7 | import einops
8 | import numpy as np
9 | import datetime
10 | import torchvision
11 |
12 | import safetensors.torch as sf
13 | from PIL import Image
14 |
15 |
16 | def min_resize(x, m):
17 | if x.shape[0] < x.shape[1]:
18 | s0 = m
19 | s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
20 | else:
21 | s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
22 | s1 = m
23 | new_max = max(s1, s0)
24 | raw_max = max(x.shape[0], x.shape[1])
25 | if new_max < raw_max:
26 | interpolation = cv2.INTER_AREA
27 | else:
28 | interpolation = cv2.INTER_LANCZOS4
29 | y = cv2.resize(x, (s1, s0), interpolation=interpolation)
30 | return y
31 |
32 |
33 | def d_resize(x, y):
34 | H, W, C = y.shape
35 | new_min = min(H, W)
36 | raw_min = min(x.shape[0], x.shape[1])
37 | if new_min < raw_min:
38 | interpolation = cv2.INTER_AREA
39 | else:
40 | interpolation = cv2.INTER_LANCZOS4
41 | y = cv2.resize(x, (W, H), interpolation=interpolation)
42 | return y
43 |
44 |
45 | def resize_and_center_crop(image, target_width, target_height):
46 | if target_height == image.shape[0] and target_width == image.shape[1]:
47 | return image
48 |
49 | pil_image = Image.fromarray(image)
50 | original_width, original_height = pil_image.size
51 | scale_factor = max(target_width / original_width, target_height / original_height)
52 | resized_width = int(round(original_width * scale_factor))
53 | resized_height = int(round(original_height * scale_factor))
54 | resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
55 | left = (resized_width - target_width) / 2
56 | top = (resized_height - target_height) / 2
57 | right = (resized_width + target_width) / 2
58 | bottom = (resized_height + target_height) / 2
59 | cropped_image = resized_image.crop((left, top, right, bottom))
60 | return np.array(cropped_image)
61 |
62 |
63 | def resize_and_center_crop_pytorch(image, target_width, target_height):
64 | B, C, H, W = image.shape
65 |
66 | if H == target_height and W == target_width:
67 | return image
68 |
69 | scale_factor = max(target_width / W, target_height / H)
70 | resized_width = int(round(W * scale_factor))
71 | resized_height = int(round(H * scale_factor))
72 |
73 | resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
74 |
75 | top = (resized_height - target_height) // 2
76 | left = (resized_width - target_width) // 2
77 | cropped = resized[:, :, top:top + target_height, left:left + target_width]
78 |
79 | return cropped
80 |
81 |
82 | def resize_without_crop(image, target_width, target_height):
83 | if target_height == image.shape[0] and target_width == image.shape[1]:
84 | return image
85 |
86 | pil_image = Image.fromarray(image)
87 | resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
88 | return np.array(resized_image)
89 |
90 |
91 | def just_crop(image, w, h):
92 | if h == image.shape[0] and w == image.shape[1]:
93 | return image
94 |
95 | original_height, original_width = image.shape[:2]
96 | k = min(original_height / h, original_width / w)
97 | new_width = int(round(w * k))
98 | new_height = int(round(h * k))
99 | x_start = (original_width - new_width) // 2
100 | y_start = (original_height - new_height) // 2
101 | cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
102 | return cropped_image
103 |
104 |
105 | def write_to_json(data, file_path):
106 | temp_file_path = file_path + ".tmp"
107 | with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
108 | json.dump(data, temp_file, indent=4)
109 | os.replace(temp_file_path, file_path)
110 | return
111 |
112 |
113 | def read_from_json(file_path):
114 | with open(file_path, 'rt', encoding='utf-8') as file:
115 | data = json.load(file)
116 | return data
117 |
118 |
119 | def get_active_parameters(m):
120 | return {k: v for k, v in m.named_parameters() if v.requires_grad}
121 |
122 |
123 | def cast_training_params(m, dtype=torch.float32):
124 | result = {}
125 | for n, param in m.named_parameters():
126 | if param.requires_grad:
127 | param.data = param.to(dtype)
128 | result[n] = param
129 | return result
130 |
131 |
132 | def separate_lora_AB(parameters, B_patterns=None):
133 | parameters_normal = {}
134 | parameters_B = {}
135 |
136 | if B_patterns is None:
137 | B_patterns = ['.lora_B.', '__zero__']
138 |
139 | for k, v in parameters.items():
140 | if any(B_pattern in k for B_pattern in B_patterns):
141 | parameters_B[k] = v
142 | else:
143 | parameters_normal[k] = v
144 |
145 | return parameters_normal, parameters_B
146 |
147 |
148 | def set_attr_recursive(obj, attr, value):
149 | attrs = attr.split(".")
150 | for name in attrs[:-1]:
151 | obj = getattr(obj, name)
152 | setattr(obj, attrs[-1], value)
153 | return
154 |
155 |
156 | def print_tensor_list_size(tensors):
157 | total_size = 0
158 | total_elements = 0
159 |
160 | if isinstance(tensors, dict):
161 | tensors = tensors.values()
162 |
163 | for tensor in tensors:
164 | total_size += tensor.nelement() * tensor.element_size()
165 | total_elements += tensor.nelement()
166 |
167 | total_size_MB = total_size / (1024 ** 2)
168 | total_elements_B = total_elements / 1e9
169 |
170 | print(f"Total number of tensors: {len(tensors)}")
171 | print(f"Total size of tensors: {total_size_MB:.2f} MB")
172 | print(f"Total number of parameters: {total_elements_B:.3f} billion")
173 | return
174 |
175 |
176 | @torch.no_grad()
177 | def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
178 | batch_size = a.size(0)
179 |
180 | if b is None:
181 | b = torch.zeros_like(a)
182 |
183 | if mask_a is None:
184 | mask_a = torch.rand(batch_size) < probability_a
185 |
186 | mask_a = mask_a.to(a.device)
187 | mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
188 | result = torch.where(mask_a, a, b)
189 | return result
190 |
191 |
192 | @torch.no_grad()
193 | def zero_module(module):
194 | for p in module.parameters():
195 | p.detach().zero_()
196 | return module
197 |
198 |
199 | @torch.no_grad()
200 | def supress_lower_channels(m, k, alpha=0.01):
201 | data = m.weight.data.clone()
202 |
203 | assert int(data.shape[1]) >= k
204 |
205 | data[:, :k] = data[:, :k] * alpha
206 | m.weight.data = data.contiguous().clone()
207 | return m
208 |
209 |
210 | def freeze_module(m):
211 | if not hasattr(m, '_forward_inside_frozen_module'):
212 | m._forward_inside_frozen_module = m.forward
213 | m.requires_grad_(False)
214 | m.forward = torch.no_grad()(m.forward)
215 | return m
216 |
217 |
218 | def get_latest_safetensors(folder_path):
219 | safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
220 |
221 | if not safetensors_files:
222 | raise ValueError('No file to resume!')
223 |
224 | latest_file = max(safetensors_files, key=os.path.getmtime)
225 | latest_file = os.path.abspath(os.path.realpath(latest_file))
226 | return latest_file
227 |
228 |
229 | def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
230 | tags = tags_str.split(', ')
231 | tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
232 | prompt = ', '.join(tags)
233 | return prompt
234 |
235 |
236 | def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
237 | numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
238 | if round_to_int:
239 | numbers = np.round(numbers).astype(int)
240 | return numbers.tolist()
241 |
242 |
243 | def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
244 | edges = np.linspace(0, 1, n + 1)
245 | points = np.random.uniform(edges[:-1], edges[1:])
246 | numbers = inclusive + (exclusive - inclusive) * points
247 | if round_to_int:
248 | numbers = np.round(numbers).astype(int)
249 | return numbers.tolist()
250 |
251 |
252 | def soft_append_bcthw(history, current, overlap=0):
253 | if overlap <= 0:
254 | return torch.cat([history, current], dim=2)
255 |
256 | assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
257 | assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
258 |
259 | weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
260 | blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
261 | output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
262 |
263 | return output.to(history)
264 |
265 |
266 | def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
267 | b, c, t, h, w = x.shape
268 |
269 | per_row = b
270 | for p in [6, 5, 4, 3, 2]:
271 | if b % p == 0:
272 | per_row = p
273 | break
274 |
275 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
276 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
277 | x = x.detach().cpu().to(torch.uint8)
278 | x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
279 | torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})
280 | return x
281 |
282 |
283 | def save_bcthw_as_png(x, output_filename):
284 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
285 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
286 | x = x.detach().cpu().to(torch.uint8)
287 | x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
288 | torchvision.io.write_png(x, output_filename)
289 | return output_filename
290 |
291 |
292 | def save_bchw_as_png(x, output_filename):
293 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
294 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
295 | x = x.detach().cpu().to(torch.uint8)
296 | x = einops.rearrange(x, 'b c h w -> c h (b w)')
297 | torchvision.io.write_png(x, output_filename)
298 | return output_filename
299 |
300 |
301 | def add_tensors_with_padding(tensor1, tensor2):
302 | if tensor1.shape == tensor2.shape:
303 | return tensor1 + tensor2
304 |
305 | shape1 = tensor1.shape
306 | shape2 = tensor2.shape
307 |
308 | new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
309 |
310 | padded_tensor1 = torch.zeros(new_shape)
311 | padded_tensor2 = torch.zeros(new_shape)
312 |
313 | padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
314 | padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
315 |
316 | result = padded_tensor1 + padded_tensor2
317 | return result
318 |
319 |
320 | def print_free_mem():
321 | torch.cuda.empty_cache()
322 | free_mem, total_mem = torch.cuda.mem_get_info(0)
323 | free_mem_mb = free_mem / (1024 ** 2)
324 | total_mem_mb = total_mem / (1024 ** 2)
325 | print(f"Free memory: {free_mem_mb:.2f} MB")
326 | print(f"Total memory: {total_mem_mb:.2f} MB")
327 | return
328 |
329 |
330 | def print_gpu_parameters(device, state_dict, log_count=1):
331 | summary = {"device": device, "keys_count": len(state_dict)}
332 |
333 | logged_params = {}
334 | for i, (key, tensor) in enumerate(state_dict.items()):
335 | if i >= log_count:
336 | break
337 | logged_params[key] = tensor.flatten()[:3].tolist()
338 |
339 | summary["params"] = logged_params
340 |
341 | print(str(summary))
342 | return
343 |
344 |
345 | def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
346 | from PIL import Image, ImageDraw, ImageFont
347 |
348 | txt = Image.new("RGB", (width, height), color="white")
349 | draw = ImageDraw.Draw(txt)
350 | font = ImageFont.truetype(font_path, size=size)
351 |
352 | if text == '':
353 | return np.array(txt)
354 |
355 | # Split text into lines that fit within the image width
356 | lines = []
357 | words = text.split()
358 | current_line = words[0]
359 |
360 | for word in words[1:]:
361 | line_with_word = f"{current_line} {word}"
362 | if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
363 | current_line = line_with_word
364 | else:
365 | lines.append(current_line)
366 | current_line = word
367 |
368 | lines.append(current_line)
369 |
370 | # Draw the text line by line
371 | y = 0
372 | line_height = draw.textbbox((0, 0), "A", font=font)[3]
373 |
374 | for line in lines:
375 | if y + line_height > height:
376 | break # stop drawing if the next line will be outside the image
377 | draw.text((0, y), line, fill="black", font=font)
378 | y += line_height
379 |
380 | return np.array(txt)
381 |
382 |
383 | def blue_mark(x):
384 | x = x.copy()
385 | c = x[:, :, 2]
386 | b = cv2.blur(c, (9, 9))
387 | x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
388 | return x
389 |
390 |
391 | def green_mark(x):
392 | x = x.copy()
393 | x[:, :, 2] = -1
394 | x[:, :, 0] = -1
395 | return x
396 |
397 |
398 | def frame_mark(x):
399 | x = x.copy()
400 | x[:64] = -1
401 | x[-64:] = -1
402 | x[:, :8] = 1
403 | x[:, -8:] = 1
404 | return x
405 |
406 |
407 | @torch.inference_mode()
408 | def pytorch2numpy(imgs):
409 | results = []
410 | for x in imgs:
411 | y = x.movedim(0, -1)
412 | y = y * 127.5 + 127.5
413 | y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
414 | results.append(y)
415 | return results
416 |
417 |
418 | @torch.inference_mode()
419 | def numpy2pytorch(imgs):
420 | h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
421 | h = h.movedim(-1, 1)
422 | return h
423 |
424 |
425 | @torch.no_grad()
426 | def duplicate_prefix_to_suffix(x, count, zero_out=False):
427 | if zero_out:
428 | return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
429 | else:
430 | return torch.cat([x, x[:count]], dim=0)
431 |
432 |
433 | def weighted_mse(a, b, weight):
434 | return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
435 |
436 |
437 | def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
438 | x = (x - x_min) / (x_max - x_min)
439 | x = max(0.0, min(x, 1.0))
440 | x = x ** sigma
441 | return y_min + x * (y_max - y_min)
442 |
443 |
444 | def expand_to_dims(x, target_dims):
445 | return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
446 |
447 |
448 | def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
449 | if tensor is None:
450 | return None
451 |
452 | first_dim = tensor.shape[0]
453 |
454 | if first_dim == batch_size:
455 | return tensor
456 |
457 | if batch_size % first_dim != 0:
458 | raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
459 |
460 | repeat_times = batch_size // first_dim
461 |
462 | return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
463 |
464 |
465 | def dim5(x):
466 | return expand_to_dims(x, 5)
467 |
468 |
469 | def dim4(x):
470 | return expand_to_dims(x, 4)
471 |
472 |
473 | def dim3(x):
474 | return expand_to_dims(x, 3)
475 |
476 |
477 | def crop_or_pad_yield_mask(x, length):
478 | B, F, C = x.shape
479 | device = x.device
480 | dtype = x.dtype
481 |
482 | if F < length:
483 | y = torch.zeros((B, length, C), dtype=dtype, device=device)
484 | mask = torch.zeros((B, length), dtype=torch.bool, device=device)
485 | y[:, :F, :] = x
486 | mask[:, :F] = True
487 | return y, mask
488 |
489 | return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
490 |
491 |
492 | def extend_dim(x, dim, minimal_length, zero_pad=False):
493 | original_length = int(x.shape[dim])
494 |
495 | if original_length >= minimal_length:
496 | return x
497 |
498 | if zero_pad:
499 | padding_shape = list(x.shape)
500 | padding_shape[dim] = minimal_length - original_length
501 | padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
502 | else:
503 | idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
504 | last_element = x[idx]
505 | padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
506 |
507 | return torch.cat([x, padding], dim=dim)
508 |
509 |
510 | def lazy_positional_encoding(t, repeats=None):
511 | if not isinstance(t, list):
512 | t = [t]
513 |
514 | from diffusers.models.embeddings import get_timestep_embedding
515 |
516 | te = torch.tensor(t)
517 | te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
518 |
519 | if repeats is None:
520 | return te
521 |
522 | te = te[:, None, :].expand(-1, repeats, -1)
523 |
524 | return te
525 |
526 |
527 | def state_dict_offset_merge(A, B, C=None):
528 | result = {}
529 | keys = A.keys()
530 |
531 | for key in keys:
532 | A_value = A[key]
533 | B_value = B[key].to(A_value)
534 |
535 | if C is None:
536 | result[key] = A_value + B_value
537 | else:
538 | C_value = C[key].to(A_value)
539 | result[key] = A_value + B_value - C_value
540 |
541 | return result
542 |
543 |
544 | def state_dict_weighted_merge(state_dicts, weights):
545 | if len(state_dicts) != len(weights):
546 | raise ValueError("Number of state dictionaries must match number of weights")
547 |
548 | if not state_dicts:
549 | return {}
550 |
551 | total_weight = sum(weights)
552 |
553 | if total_weight == 0:
554 | raise ValueError("Sum of weights cannot be zero")
555 |
556 | normalized_weights = [w / total_weight for w in weights]
557 |
558 | keys = state_dicts[0].keys()
559 | result = {}
560 |
561 | for key in keys:
562 | result[key] = state_dicts[0][key] * normalized_weights[0]
563 |
564 | for i in range(1, len(state_dicts)):
565 | state_dict_value = state_dicts[i][key].to(result[key])
566 | result[key] += state_dict_value * normalized_weights[i]
567 |
568 | return result
569 |
570 |
571 | def group_files_by_folder(all_files):
572 | grouped_files = {}
573 |
574 | for file in all_files:
575 | folder_name = os.path.basename(os.path.dirname(file))
576 | if folder_name not in grouped_files:
577 | grouped_files[folder_name] = []
578 | grouped_files[folder_name].append(file)
579 |
580 | list_of_lists = list(grouped_files.values())
581 | return list_of_lists
582 |
583 |
584 | def generate_timestamp():
585 | now = datetime.datetime.now()
586 | timestamp = now.strftime('%y%m%d_%H%M%S')
587 | milliseconds = f"{int(now.microsecond / 1000):03d}"
588 | random_number = random.randint(0, 9999)
589 | return f"{timestamp}_{milliseconds}_{random_number}"
590 |
591 |
592 | def write_PIL_image_with_png_info(image, metadata, path):
593 | from PIL.PngImagePlugin import PngInfo
594 |
595 | png_info = PngInfo()
596 | for key, value in metadata.items():
597 | png_info.add_text(key, value)
598 |
599 | image.save(path, "PNG", pnginfo=png_info)
600 | return image
601 |
602 |
603 | def torch_safe_save(content, path):
604 | torch.save(content, path + '_tmp')
605 | os.replace(path + '_tmp', path)
606 | return path
607 |
608 |
609 | def move_optimizer_to_device(optimizer, device):
610 | for state in optimizer.state.values():
611 | for k, v in state.items():
612 | if isinstance(v, torch.Tensor):
613 | state[k] = v.to(device)
614 |
--------------------------------------------------------------------------------
/demo_gradio_f1.py:
--------------------------------------------------------------------------------
1 | from diffusers_helper.hf_login import login
2 |
3 | import os
4 |
5 | os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
6 | os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
7 |
8 | import gradio as gr
9 | import torch
10 | import traceback
11 | import einops
12 | import safetensors.torch as sf
13 | import numpy as np
14 | import argparse
15 | import math
16 | import gc
17 | import time
18 |
19 | from PIL import Image
20 | from diffusers import AutoencoderKLHunyuanVideo
21 | from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
22 | from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
23 | from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp
24 | from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
25 | from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
26 | from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete
27 | from diffusers_helper.thread_utils import AsyncStream, async_run
28 | from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
29 | from transformers import SiglipImageProcessor, SiglipVisionModel
30 | from diffusers_helper.clip_vision import hf_clip_vision_encode
31 | from diffusers_helper.bucket_tools import find_nearest_bucket
32 | from utils.lora_utils import merge_lora_to_state_dict
33 |
34 |
35 | parser = argparse.ArgumentParser()
36 | parser.add_argument('--share', action='store_true')
37 | parser.add_argument("--server", type=str, default='0.0.0.0')
38 | parser.add_argument("--port", type=int, required=False)
39 | parser.add_argument("--inbrowser", action='store_true')
40 | parser.add_argument("--output_dir", type=str, default='./outputs')
41 | args = parser.parse_args()
42 |
43 | # for win desktop probably use --server 127.0.0.1 --inbrowser
44 | # For linux server probably use --server 127.0.0.1 or do not use any cmd flags
45 |
46 | print(args)
47 |
48 | if torch.cuda.is_available():
49 | free_mem_gb = get_cuda_free_memory_gb(gpu)
50 | else:
51 | free_mem_gb = torch.mps.recommended_max_memory() / 1024 / 1024 / 1024
52 |
53 | high_vram = free_mem_gb > 60
54 |
55 | print(f'Free VRAM {free_mem_gb} GB')
56 | print(f'High-VRAM Mode: {high_vram}')
57 |
58 | text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
59 | text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
60 | tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
61 | tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
62 | vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
63 |
64 | feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
65 | image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
66 |
67 | transformer = None # load later
68 | transformer_dtype = torch.bfloat16
69 | previous_lora_file = None
70 | previous_lora_multiplier = None
71 |
72 | vae.eval()
73 | text_encoder.eval()
74 | text_encoder_2.eval()
75 | image_encoder.eval()
76 |
77 | if not high_vram:
78 | vae.enable_slicing()
79 | vae.enable_tiling()
80 |
81 | vae.to(dtype=torch.float16)
82 | image_encoder.to(dtype=torch.float16)
83 | text_encoder.to(dtype=torch.float16)
84 | text_encoder_2.to(dtype=torch.float16)
85 |
86 | vae.requires_grad_(False)
87 | text_encoder.requires_grad_(False)
88 | text_encoder_2.requires_grad_(False)
89 | image_encoder.requires_grad_(False)
90 |
91 | if not high_vram:
92 | # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
93 | DynamicSwapInstaller.install_model(text_encoder, device=gpu)
94 | else:
95 | text_encoder.to(gpu)
96 | text_encoder_2.to(gpu)
97 | image_encoder.to(gpu)
98 | vae.to(gpu)
99 |
100 | stream = AsyncStream()
101 |
102 | outputs_folder = args.output_dir
103 | os.makedirs(outputs_folder, exist_ok=True)
104 |
105 |
106 | @torch.no_grad()
107 | def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier):
108 | global transformer, previous_lora_file, previous_lora_multiplier
109 |
110 | model_changed = transformer is None or (
111 | lora_file != previous_lora_file
112 | or lora_multiplier != previous_lora_multiplier
113 | )
114 |
115 | total_latent_sections = (total_second_length * 24) / (latent_window_size * 4)
116 | total_latent_sections = int(max(round(total_latent_sections), 1))
117 |
118 | job_id = generate_timestamp()
119 |
120 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
121 |
122 | try:
123 | # Clean GPU
124 | if not high_vram:
125 | unload_complete_models(
126 | text_encoder, text_encoder_2, image_encoder, vae, transformer
127 | )
128 |
129 | # Text encoding
130 |
131 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
132 |
133 | if not high_vram:
134 | fake_diffusers_current_device(text_encoder, gpu) # since we only encode one text - that is one model move and one encode, offload is same time consumption since it is also one load and one encode.
135 | load_model_as_complete(text_encoder_2, target_device=gpu)
136 |
137 | llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
138 |
139 | if cfg == 1:
140 | llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
141 | else:
142 | llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
143 |
144 | llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
145 | llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
146 |
147 | # Processing input image
148 |
149 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
150 |
151 | H, W, C = input_image.shape
152 | height, width = find_nearest_bucket(H, W, resolution=resolution)
153 | input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
154 |
155 | Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
156 |
157 | input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
158 | input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
159 |
160 | # VAE encoding
161 |
162 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
163 |
164 | if not high_vram:
165 | load_model_as_complete(vae, target_device=gpu)
166 |
167 | start_latent = vae_encode(input_image_pt, vae)
168 |
169 | # CLIP Vision
170 |
171 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
172 |
173 | if not high_vram:
174 | load_model_as_complete(image_encoder, target_device=gpu)
175 |
176 | image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
177 | image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
178 |
179 | # Dtype
180 |
181 | llama_vec = llama_vec.to(transformer_dtype)
182 | llama_vec_n = llama_vec_n.to(transformer_dtype)
183 | clip_l_pooler = clip_l_pooler.to(transformer_dtype)
184 | clip_l_pooler_n = clip_l_pooler_n.to(transformer_dtype)
185 | image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer_dtype)
186 |
187 | # Load transformer model
188 | if model_changed:
189 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Loading transformer ...'))))
190 |
191 | transformer = None
192 | time.sleep(1.0) # wait for the previous model to be unloaded
193 | torch.cuda.empty_cache()
194 | gc.collect()
195 |
196 | previous_lora_file = lora_file
197 | previous_lora_multiplier = lora_multiplier
198 |
199 | transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503', torch_dtype=torch.bfloat16).cpu()
200 | transformer.eval()
201 | transformer.high_quality_fp32_output_for_inference = True
202 | print('transformer.high_quality_fp32_output_for_inference = True')
203 |
204 | transformer.to(dtype=torch.bfloat16)
205 | transformer.requires_grad_(False)
206 |
207 | if lora_file is not None:
208 | state_dict = transformer.state_dict()
209 | print(f"Merging LoRA file {os.path.basename(lora_file)} ...")
210 | state_dict = merge_lora_to_state_dict(state_dict, lora_file, lora_multiplier, device=gpu)
211 | gc.collect()
212 | info = transformer.load_state_dict(state_dict, strict=True, assign=True)
213 | print(f"LoRA applied: {info}")
214 |
215 | if not high_vram:
216 | DynamicSwapInstaller.install_model(transformer, device=gpu)
217 | else:
218 | transformer.to(gpu)
219 |
220 | # Sampling
221 |
222 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
223 |
224 | rnd = torch.Generator("cpu").manual_seed(seed)
225 |
226 | history_latents = torch.zeros(size=(1, 16, 16 + 2 + 1, height // 8, width // 8), dtype=torch.float32).cpu()
227 | history_pixels = None
228 |
229 | history_latents = torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
230 | total_generated_latent_frames = 1
231 |
232 | for section_index in range(total_latent_sections):
233 | if stream.input_queue.top() == 'end':
234 | stream.output_queue.push(('end', None))
235 | return
236 |
237 | print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
238 |
239 | if not high_vram:
240 | unload_complete_models()
241 | move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
242 |
243 | if use_teacache:
244 | transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
245 | else:
246 | transformer.initialize_teacache(enable_teacache=False)
247 |
248 | def callback(d):
249 | preview = d['denoised']
250 | preview = vae_decode_fake(preview)
251 |
252 | preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
253 | preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
254 |
255 | if stream.input_queue.top() == 'end':
256 | stream.output_queue.push(('end', None))
257 | raise KeyboardInterrupt('User ends the task.')
258 |
259 | current_step = d['i'] + 1
260 | percentage = int(100.0 * current_step / steps)
261 | hint = f'Sampling {current_step}/{steps}'
262 | desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 24) :.2f} seconds (FPS-24). The video is being extended now ...'
263 | stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
264 | return
265 |
266 | indices = torch.arange(0, sum([1, 16, 2, 1, latent_window_size])).unsqueeze(0)
267 | clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
268 | clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
269 |
270 | clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2)
271 | clean_latents = torch.cat([start_latent.to(history_latents), clean_latents_1x], dim=2)
272 |
273 | generated_latents = sample_hunyuan(
274 | transformer=transformer,
275 | sampler='unipc',
276 | width=width,
277 | height=height,
278 | frames=latent_window_size * 4 - 3,
279 | real_guidance_scale=cfg,
280 | distilled_guidance_scale=gs,
281 | guidance_rescale=rs,
282 | # shift=3.0,
283 | num_inference_steps=steps,
284 | generator=rnd,
285 | prompt_embeds=llama_vec,
286 | prompt_embeds_mask=llama_attention_mask,
287 | prompt_poolers=clip_l_pooler,
288 | negative_prompt_embeds=llama_vec_n,
289 | negative_prompt_embeds_mask=llama_attention_mask_n,
290 | negative_prompt_poolers=clip_l_pooler_n,
291 | device=gpu,
292 | dtype=transformer.dtype,
293 | image_embeddings=image_encoder_last_hidden_state,
294 | latent_indices=latent_indices,
295 | clean_latents=clean_latents,
296 | clean_latent_indices=clean_latent_indices,
297 | clean_latents_2x=clean_latents_2x,
298 | clean_latent_2x_indices=clean_latent_2x_indices,
299 | clean_latents_4x=clean_latents_4x,
300 | clean_latent_4x_indices=clean_latent_4x_indices,
301 | callback=callback,
302 | )
303 |
304 | total_generated_latent_frames += int(generated_latents.shape[2])
305 | history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
306 |
307 | if not high_vram:
308 | offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
309 | load_model_as_complete(vae, target_device=gpu)
310 |
311 | real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
312 |
313 | if history_pixels is None:
314 | history_pixels = vae_decode(real_history_latents, vae).cpu()
315 | else:
316 | section_latent_frames = latent_window_size * 2
317 | overlapped_frames = latent_window_size * 4 - 3
318 |
319 | current_pixels = vae_decode(real_history_latents[:, :, -section_latent_frames:], vae).cpu()
320 | history_pixels = soft_append_bcthw(history_pixels, current_pixels, overlapped_frames)
321 |
322 | if not high_vram:
323 | unload_complete_models()
324 |
325 | output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
326 |
327 | save_bcthw_as_mp4(history_pixels, output_filename, fps=24, crf=mp4_crf)
328 |
329 | print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
330 |
331 | stream.output_queue.push(('file', output_filename))
332 | except:
333 | traceback.print_exc()
334 |
335 | if not high_vram:
336 | unload_complete_models(
337 | text_encoder, text_encoder_2, image_encoder, vae, transformer
338 | )
339 |
340 | stream.output_queue.push(('end', None))
341 | return
342 |
343 |
344 | def process(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier):
345 | global stream
346 | assert input_image is not None, 'No input image!'
347 |
348 | yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
349 |
350 | stream = AsyncStream()
351 |
352 | async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier)
353 |
354 | output_filename = None
355 |
356 | while True:
357 | flag, data = stream.output_queue.next()
358 |
359 | if flag == 'file':
360 | output_filename = data
361 | yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
362 |
363 | if flag == 'progress':
364 | preview, desc, html = data
365 | yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
366 |
367 | if flag == 'end':
368 | yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
369 | break
370 |
371 |
372 | def end_process():
373 | stream.input_queue.push('end')
374 |
375 |
376 | quick_prompts = [
377 | 'The girl dances gracefully, with clear movements, full of charm.',
378 | 'A character doing some simple body movements.',
379 | ]
380 | quick_prompts = [[x] for x in quick_prompts]
381 |
382 |
383 | css = make_progress_bar_css()
384 | block = gr.Blocks(css=css).queue()
385 | with block:
386 | gr.Markdown('# FramePack-F1')
387 | with gr.Row():
388 | with gr.Column():
389 | input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
390 | resolution = gr.Slider(label="Resolution", minimum=240, maximum=720, value=416, step=16)
391 | prompt = gr.Textbox(label="Prompt", value='')
392 | example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
393 | example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
394 |
395 | with gr.Row():
396 | start_button = gr.Button(value="Start Generation")
397 | end_button = gr.Button(value="End Generation", interactive=False)
398 |
399 | with gr.Group():
400 | use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
401 |
402 | n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=False) # Not used
403 | seed = gr.Number(label="Seed", value=31337, precision=0)
404 |
405 | total_second_length = gr.Slider(label="Total Video Length (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
406 | latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False) # Should not change
407 | steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Changing this value is not recommended.')
408 |
409 | cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False) # Should not change
410 | gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.01, info='Changing this value is not recommended.')
411 | rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
412 |
413 | # This is only used when high_vram is False
414 | gpu_memory_preservation = gr.Slider(label="GPU Inference Preserved Memory (GB) (larger means slower)", minimum=6, maximum=128, value=6, step=0.1, info="Set this number to a larger value if you encounter OOM. Larger value causes slower speed.", visible=not high_vram)
415 |
416 | mp4_crf = gr.Slider(label="MP4 Compression", minimum=0, maximum=100, value=16, step=1, info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs. ")
417 |
418 | with gr.Group():
419 | lora_file = gr.File(label="LoRA File", file_count="single", type="filepath")
420 | lora_multiplier = gr.Slider(label="LoRA Multiplier", minimum=0.0, maximum=1.0, value=0.8, step=0.1)
421 |
422 | with gr.Column():
423 | preview_image = gr.Image(label="Next Latents", height=200, visible=False)
424 | result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
425 | progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
426 | progress_bar = gr.HTML('', elem_classes='no-generating-animation')
427 |
428 | gr.HTML('')
429 |
430 | ips = [input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier]
431 | start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
432 | end_button.click(fn=end_process)
433 |
434 |
435 | block.launch(
436 | server_name=args.server,
437 | server_port=args.port,
438 | share=args.share,
439 | inbrowser=args.inbrowser,
440 | allowed_paths=[outputs_folder],
441 | )
442 |
--------------------------------------------------------------------------------
/demo_gradio.py:
--------------------------------------------------------------------------------
1 | from diffusers_helper.hf_login import login
2 |
3 | import os
4 |
5 | os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
6 | os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
7 |
8 | import gradio as gr
9 | import torch
10 | import traceback
11 | import einops
12 | import safetensors.torch as sf
13 | import numpy as np
14 | import argparse
15 | import math
16 | import gc
17 | import time
18 |
19 | from PIL import Image
20 | from diffusers import AutoencoderKLHunyuanVideo
21 | from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
22 | from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
23 | from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp
24 | from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
25 | from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
26 | from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete
27 | from diffusers_helper.thread_utils import AsyncStream, async_run
28 | from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
29 | from transformers import SiglipImageProcessor, SiglipVisionModel
30 | from diffusers_helper.clip_vision import hf_clip_vision_encode
31 | from diffusers_helper.bucket_tools import find_nearest_bucket
32 | from utils.lora_utils import merge_lora_to_state_dict
33 |
34 |
35 | parser = argparse.ArgumentParser()
36 | parser.add_argument('--share', action='store_true')
37 | parser.add_argument("--server", type=str, default='0.0.0.0')
38 | parser.add_argument("--port", type=int, required=False)
39 | parser.add_argument("--inbrowser", action='store_true')
40 | parser.add_argument("--output_dir", type=str, default='./outputs')
41 | args = parser.parse_args()
42 |
43 | # for win desktop probably use --server 127.0.0.1 --inbrowser
44 | # For linux server probably use --server 127.0.0.1 or do not use any cmd flags
45 |
46 | print(args)
47 |
48 | if torch.cuda.is_available():
49 | free_mem_gb = get_cuda_free_memory_gb(gpu)
50 | else:
51 | free_mem_gb = torch.mps.recommended_max_memory() / 1024 / 1024 / 1024
52 |
53 | high_vram = free_mem_gb > 60
54 | print(f'Free VRAM {free_mem_gb} GB')
55 | print(f'High-VRAM Mode: {high_vram}')
56 |
57 | text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
58 | text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
59 | tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
60 | tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
61 | vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
62 |
63 | feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
64 | image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
65 |
66 | transformer = None # load later
67 | transformer_dtype = torch.bfloat16
68 | previous_lora_file = None
69 | previous_lora_multiplier = None
70 |
71 | vae.eval()
72 | text_encoder.eval()
73 | text_encoder_2.eval()
74 | image_encoder.eval()
75 |
76 | if not high_vram:
77 | vae.enable_slicing()
78 | vae.enable_tiling()
79 |
80 | vae.to(dtype=torch.float16)
81 | image_encoder.to(dtype=torch.float16)
82 | text_encoder.to(dtype=torch.float16)
83 | text_encoder_2.to(dtype=torch.float16)
84 |
85 | vae.requires_grad_(False)
86 | text_encoder.requires_grad_(False)
87 | text_encoder_2.requires_grad_(False)
88 | image_encoder.requires_grad_(False)
89 |
90 | if not high_vram:
91 | # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
92 | DynamicSwapInstaller.install_model(text_encoder, device=gpu)
93 | else:
94 | text_encoder.to(gpu)
95 | text_encoder_2.to(gpu)
96 | image_encoder.to(gpu)
97 | vae.to(gpu)
98 |
99 | stream = AsyncStream()
100 |
101 | outputs_folder = args.output_dir
102 | os.makedirs(outputs_folder, exist_ok=True)
103 |
104 |
105 | @torch.no_grad()
106 | def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier):
107 | global transformer, previous_lora_file, previous_lora_multiplier
108 |
109 | model_changed = transformer is None or (
110 | lora_file != previous_lora_file
111 | or lora_multiplier != previous_lora_multiplier
112 | )
113 |
114 | total_latent_sections = (total_second_length * 24) / (latent_window_size * 4)
115 | total_latent_sections = int(max(round(total_latent_sections), 1))
116 |
117 | job_id = generate_timestamp()
118 |
119 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
120 |
121 | try:
122 | # Clean GPU
123 | if not high_vram:
124 | unload_complete_models(
125 | text_encoder, text_encoder_2, image_encoder, vae, transformer
126 | )
127 |
128 | # Text encoding
129 |
130 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
131 |
132 | if not high_vram:
133 | fake_diffusers_current_device(text_encoder, gpu)
134 | load_model_as_complete(text_encoder_2, target_device=gpu)
135 |
136 | llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
137 |
138 | if cfg == 1:
139 | llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
140 | else:
141 | llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
142 |
143 | llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
144 | llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
145 |
146 | # Processing input image
147 |
148 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
149 |
150 | H, W, C = input_image.shape
151 | height, width = find_nearest_bucket(H, W, resolution=resolution)
152 | input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
153 |
154 | Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
155 |
156 | input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
157 | input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
158 |
159 | # VAE encoding
160 |
161 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
162 |
163 | if not high_vram:
164 | load_model_as_complete(vae, target_device=gpu)
165 |
166 | start_latent = vae_encode(input_image_pt, vae)
167 |
168 | # CLIP Vision
169 |
170 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
171 |
172 | if not high_vram:
173 | load_model_as_complete(image_encoder, target_device=gpu)
174 |
175 | image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
176 | image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
177 |
178 | # Dtype
179 |
180 | llama_vec = llama_vec.to(transformer_dtype)
181 | llama_vec_n = llama_vec_n.to(transformer_dtype)
182 | clip_l_pooler = clip_l_pooler.to(transformer_dtype)
183 | clip_l_pooler_n = clip_l_pooler_n.to(transformer_dtype)
184 | image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer_dtype)
185 |
186 | # Load transformer model
187 | if model_changed:
188 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Loading transformer ...'))))
189 |
190 | transformer = None
191 | time.sleep(1.0) # wait for the previous model to be unloaded
192 | torch.cuda.empty_cache()
193 | gc.collect()
194 |
195 | previous_lora_file = lora_file
196 | previous_lora_multiplier = lora_multiplier
197 |
198 | transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu()
199 | transformer.eval()
200 | transformer.high_quality_fp32_output_for_inference = True
201 | print('transformer.high_quality_fp32_output_for_inference = True')
202 |
203 | transformer.to(dtype=torch.bfloat16)
204 | transformer.requires_grad_(False)
205 |
206 | if lora_file is not None:
207 | state_dict = transformer.state_dict()
208 | print(f"Merging LoRA file {os.path.basename(lora_file)} ...")
209 | state_dict = merge_lora_to_state_dict(state_dict, lora_file, lora_multiplier, device=gpu)
210 | gc.collect()
211 | info = transformer.load_state_dict(state_dict, strict=True, assign=True)
212 | print(f"LoRA applied: {info}")
213 |
214 | if not high_vram:
215 | DynamicSwapInstaller.install_model(transformer, device=gpu)
216 | else:
217 | transformer.to(gpu)
218 |
219 | # Sampling
220 |
221 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
222 |
223 | rnd = torch.Generator("cpu").manual_seed(seed)
224 | num_frames = latent_window_size * 4 - 3
225 |
226 | history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32).cpu()
227 | history_pixels = None
228 | total_generated_latent_frames = 0
229 |
230 | latent_paddings = reversed(range(total_latent_sections))
231 |
232 | if total_latent_sections > 4:
233 | latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
234 |
235 | for latent_padding in latent_paddings:
236 | is_last_section = latent_padding == 0
237 | latent_padding_size = latent_padding * latent_window_size
238 |
239 | if stream.input_queue.top() == 'end':
240 | stream.output_queue.push(('end', None))
241 | return
242 |
243 | print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}')
244 |
245 | indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0)
246 | clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split([1, latent_padding_size, latent_window_size, 1, 2, 16], dim=1)
247 | clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1)
248 |
249 | clean_latents_pre = start_latent.to(history_latents)
250 | clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2)
251 | clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2)
252 |
253 | if not high_vram:
254 | unload_complete_models()
255 | move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
256 |
257 | if use_teacache:
258 | transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
259 | else:
260 | transformer.initialize_teacache(enable_teacache=False)
261 |
262 | def callback(d):
263 | preview = d['denoised']
264 | preview = vae_decode_fake(preview)
265 |
266 | preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
267 | preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
268 |
269 | if stream.input_queue.top() == 'end':
270 | stream.output_queue.push(('end', None))
271 | raise KeyboardInterrupt('User ends the task.')
272 |
273 | current_step = d['i'] + 1
274 | percentage = int(100.0 * current_step / steps)
275 | hint = f'Sampling {current_step}/{steps}'
276 | desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 24) :.2f} seconds (FPS-24). The video is being extended now ...'
277 | stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
278 | return
279 |
280 | generated_latents = sample_hunyuan(
281 | transformer=transformer,
282 | sampler='unipc',
283 | width=width,
284 | height=height,
285 | frames=num_frames,
286 | real_guidance_scale=cfg,
287 | distilled_guidance_scale=gs,
288 | guidance_rescale=rs,
289 | num_inference_steps=steps,
290 | generator=rnd,
291 | prompt_embeds=llama_vec,
292 | prompt_embeds_mask=llama_attention_mask,
293 | prompt_poolers=clip_l_pooler,
294 | negative_prompt_embeds=llama_vec_n,
295 | negative_prompt_embeds_mask=llama_attention_mask_n,
296 | negative_prompt_poolers=clip_l_pooler_n,
297 | device=gpu,
298 | dtype=transformer.dtype,
299 | image_embeddings=image_encoder_last_hidden_state,
300 | latent_indices=latent_indices,
301 | clean_latents=clean_latents,
302 | clean_latent_indices=clean_latent_indices,
303 | clean_latents_2x=clean_latents_2x,
304 | clean_latent_2x_indices=clean_latent_2x_indices,
305 | clean_latents_4x=clean_latents_4x,
306 | clean_latent_4x_indices=clean_latent_4x_indices,
307 | callback=callback,
308 | )
309 |
310 | if is_last_section:
311 | generated_latents = torch.cat([start_latent.to(generated_latents), generated_latents], dim=2)
312 |
313 | total_generated_latent_frames += int(generated_latents.shape[2])
314 | history_latents = torch.cat([generated_latents.to(history_latents), history_latents], dim=2)
315 |
316 | if not high_vram:
317 | offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
318 | load_model_as_complete(vae, target_device=gpu)
319 |
320 | real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
321 |
322 | if history_pixels is None:
323 | history_pixels = vae_decode(real_history_latents, vae).cpu()
324 | else:
325 | section_latent_frames = (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2)
326 | overlapped_frames = latent_window_size * 4 - 3
327 |
328 | current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], vae).cpu()
329 | history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
330 |
331 | if not high_vram:
332 | unload_complete_models()
333 |
334 | output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
335 |
336 | save_bcthw_as_mp4(history_pixels, output_filename, fps=24, crf=mp4_crf)
337 |
338 | print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
339 |
340 | stream.output_queue.push(('file', output_filename))
341 |
342 | if is_last_section:
343 | break
344 | except:
345 | traceback.print_exc()
346 |
347 | if not high_vram:
348 | unload_complete_models(
349 | text_encoder, text_encoder_2, image_encoder, vae, transformer
350 | )
351 |
352 | stream.output_queue.push(('end', None))
353 | return
354 |
355 |
356 | def process(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier):
357 | global stream
358 | assert input_image is not None, 'No input image!'
359 |
360 | yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
361 |
362 | stream = AsyncStream()
363 |
364 | async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier)
365 |
366 | output_filename = None
367 |
368 | while True:
369 | flag, data = stream.output_queue.next()
370 |
371 | if flag == 'file':
372 | output_filename = data
373 | yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
374 |
375 | if flag == 'progress':
376 | preview, desc, html = data
377 | yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
378 |
379 | if flag == 'end':
380 | yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
381 | break
382 |
383 |
384 | def end_process():
385 | stream.input_queue.push('end')
386 |
387 |
388 | quick_prompts = [
389 | 'The girl dances gracefully, with clear movements, full of charm.',
390 | 'A character doing some simple body movements.',
391 | ]
392 | quick_prompts = [[x] for x in quick_prompts]
393 |
394 |
395 | css = make_progress_bar_css()
396 | block = gr.Blocks(css=css).queue()
397 | with block:
398 | gr.Markdown('# FramePack')
399 | with gr.Row():
400 | with gr.Column():
401 | input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
402 | resolution = gr.Slider(label="Resolution", minimum=240, maximum=720, value=416, step=16)
403 | prompt = gr.Textbox(label="Prompt", value='')
404 | example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
405 | example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
406 |
407 | with gr.Row():
408 | start_button = gr.Button(value="Start Generation")
409 | end_button = gr.Button(value="End Generation", interactive=False)
410 |
411 | with gr.Group():
412 | use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
413 |
414 | n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=False) # Not used
415 | seed = gr.Number(label="Seed", value=31337, precision=0)
416 |
417 | total_second_length = gr.Slider(label="Total Video Length (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
418 | latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False) # Should not change
419 | steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Changing this value is not recommended.')
420 |
421 | cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False) # Should not change
422 | gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.01, info='Changing this value is not recommended.')
423 | rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
424 |
425 | # This is only used when high_vram is False
426 | gpu_memory_preservation = gr.Slider(label="GPU Inference Preserved Memory (GB) (larger means slower)", minimum=6, maximum=128, value=6, step=0.1, info="Set this number to a larger value if you encounter OOM. Larger value causes slower speed.", visible=not high_vram)
427 |
428 | mp4_crf = gr.Slider(label="MP4 Compression", minimum=0, maximum=100, value=16, step=1, info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs. ")
429 |
430 | with gr.Group():
431 | lora_file = gr.File(label="LoRA File", file_count="single", type="filepath")
432 | lora_multiplier = gr.Slider(label="LoRA Multiplier", minimum=0.0, maximum=1.0, value=0.8, step=0.1)
433 |
434 | with gr.Column():
435 | preview_image = gr.Image(label="Next Latents", height=200, visible=False)
436 | result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
437 | gr.Markdown('Note that the ending actions will be generated before the starting actions due to the inverted sampling. If the starting action is not in the video, you just need to wait, and it will be generated later.')
438 | progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
439 | progress_bar = gr.HTML('', elem_classes='no-generating-animation')
440 |
441 | gr.HTML('')
442 |
443 | ips = [input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, resolution, lora_file, lora_multiplier]
444 | start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
445 | end_button.click(fn=end_process)
446 |
447 |
448 | block.launch(
449 | server_name=args.server,
450 | server_port=args.port,
451 | share=args.share,
452 | inbrowser=args.inbrowser,
453 | allowed_paths=[outputs_folder],
454 | )
455 |
--------------------------------------------------------------------------------
/diffusers_helper/models/hunyuan_video_packed.py:
--------------------------------------------------------------------------------
1 | from typing import Any, Dict, List, Optional, Tuple, Union
2 |
3 | import torch
4 | import einops
5 | import torch.nn as nn
6 | import numpy as np
7 |
8 | from diffusers.loaders import FromOriginalModelMixin
9 | from diffusers.configuration_utils import ConfigMixin, register_to_config
10 | from diffusers.loaders import PeftAdapterMixin
11 | from diffusers.utils import logging
12 | from diffusers.models.attention import FeedForward
13 | from diffusers.models.attention_processor import Attention
14 | from diffusers.models.embeddings import TimestepEmbedding, Timesteps, PixArtAlphaTextProjection
15 | from diffusers.models.modeling_outputs import Transformer2DModelOutput
16 | from diffusers.models.modeling_utils import ModelMixin
17 | from diffusers_helper.dit_common import LayerNorm
18 | from diffusers_helper.utils import zero_module
19 |
20 |
21 | enabled_backends = []
22 |
23 | if torch.backends.cuda.flash_sdp_enabled():
24 | enabled_backends.append("flash")
25 | if torch.backends.cuda.math_sdp_enabled():
26 | enabled_backends.append("math")
27 | if torch.backends.cuda.mem_efficient_sdp_enabled():
28 | enabled_backends.append("mem_efficient")
29 | if torch.backends.cuda.cudnn_sdp_enabled():
30 | enabled_backends.append("cudnn")
31 |
32 | print("Currently enabled native sdp backends:", enabled_backends)
33 |
34 | try:
35 | # raise NotImplementedError
36 | from xformers.ops import memory_efficient_attention as xformers_attn_func
37 | print('Xformers is installed!')
38 | except:
39 | print('Xformers is not installed!')
40 | xformers_attn_func = None
41 |
42 | try:
43 | # raise NotImplementedError
44 | from flash_attn import flash_attn_varlen_func, flash_attn_func
45 | print('Flash Attn is installed!')
46 | except:
47 | print('Flash Attn is not installed!')
48 | flash_attn_varlen_func = None
49 | flash_attn_func = None
50 |
51 | try:
52 | # raise NotImplementedError
53 | from sageattention import sageattn_varlen, sageattn
54 | print('Sage Attn is installed!')
55 | except:
56 | print('Sage Attn is not installed!')
57 | sageattn_varlen = None
58 | sageattn = None
59 |
60 |
61 | logger = logging.get_logger(__name__) # pylint: disable=invalid-name
62 |
63 |
64 | def pad_for_3d_conv(x, kernel_size):
65 | b, c, t, h, w = x.shape
66 | pt, ph, pw = kernel_size
67 | pad_t = (pt - (t % pt)) % pt
68 | pad_h = (ph - (h % ph)) % ph
69 | pad_w = (pw - (w % pw)) % pw
70 | return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode='replicate')
71 |
72 |
73 | def center_down_sample_3d(x, kernel_size):
74 | # pt, ph, pw = kernel_size
75 | # cp = (pt * ph * pw) // 2
76 | # xp = einops.rearrange(x, 'b c (t pt) (h ph) (w pw) -> (pt ph pw) b c t h w', pt=pt, ph=ph, pw=pw)
77 | # xc = xp[cp]
78 | # return xc
79 | return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size)
80 |
81 |
82 | def get_cu_seqlens(text_mask, img_len):
83 | batch_size = text_mask.shape[0]
84 | text_len = text_mask.sum(dim=1)
85 | max_len = text_mask.shape[1] + img_len
86 |
87 | # Use the same device as the input tensor
88 | device = text_mask.device
89 | cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device=device)
90 |
91 | for i in range(batch_size):
92 | s = text_len[i] + img_len
93 | s1 = i * max_len + s
94 | s2 = (i + 1) * max_len
95 | cu_seqlens[2 * i + 1] = s1
96 | cu_seqlens[2 * i + 2] = s2
97 |
98 | return cu_seqlens
99 |
100 |
101 | def apply_rotary_emb_transposed(x, freqs_cis):
102 | cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
103 | x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1)
104 | x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
105 | out = x.float() * cos + x_rotated.float() * sin
106 | out = out.to(x)
107 | return out
108 |
109 |
110 | def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv):
111 | if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None:
112 | if sageattn is not None:
113 | x = sageattn(q, k, v, tensor_layout='NHD')
114 | return x
115 |
116 | if flash_attn_func is not None:
117 | x = flash_attn_func(q, k, v)
118 | return x
119 |
120 | if xformers_attn_func is not None:
121 | x = xformers_attn_func(q, k, v)
122 | return x
123 |
124 | x = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
125 | return x
126 |
127 | B, L, H, C = q.shape
128 |
129 | q = q.flatten(0, 1)
130 | k = k.flatten(0, 1)
131 | v = v.flatten(0, 1)
132 |
133 | if sageattn_varlen is not None:
134 | x = sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
135 | elif flash_attn_varlen_func is not None:
136 | x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
137 | else:
138 | raise NotImplementedError('No Attn Installed!')
139 |
140 | x = x.unflatten(0, (B, L))
141 |
142 | return x
143 |
144 |
145 | class HunyuanAttnProcessorFlashAttnDouble:
146 | def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
147 | cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
148 |
149 | query = attn.to_q(hidden_states)
150 | key = attn.to_k(hidden_states)
151 | value = attn.to_v(hidden_states)
152 |
153 | query = query.unflatten(2, (attn.heads, -1))
154 | key = key.unflatten(2, (attn.heads, -1))
155 | value = value.unflatten(2, (attn.heads, -1))
156 |
157 | query = attn.norm_q(query)
158 | key = attn.norm_k(key)
159 |
160 | query = apply_rotary_emb_transposed(query, image_rotary_emb)
161 | key = apply_rotary_emb_transposed(key, image_rotary_emb)
162 |
163 | encoder_query = attn.add_q_proj(encoder_hidden_states)
164 | encoder_key = attn.add_k_proj(encoder_hidden_states)
165 | encoder_value = attn.add_v_proj(encoder_hidden_states)
166 |
167 | encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
168 | encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
169 | encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
170 |
171 | encoder_query = attn.norm_added_q(encoder_query)
172 | encoder_key = attn.norm_added_k(encoder_key)
173 |
174 | query = torch.cat([query, encoder_query], dim=1)
175 | key = torch.cat([key, encoder_key], dim=1)
176 | value = torch.cat([value, encoder_value], dim=1)
177 |
178 | hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
179 | hidden_states = hidden_states.flatten(-2)
180 |
181 | txt_length = encoder_hidden_states.shape[1]
182 | hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
183 |
184 | hidden_states = attn.to_out[0](hidden_states)
185 | hidden_states = attn.to_out[1](hidden_states)
186 | encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
187 |
188 | return hidden_states, encoder_hidden_states
189 |
190 |
191 | class HunyuanAttnProcessorFlashAttnSingle:
192 | def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
193 | cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
194 |
195 | hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
196 |
197 | query = attn.to_q(hidden_states)
198 | key = attn.to_k(hidden_states)
199 | value = attn.to_v(hidden_states)
200 |
201 | query = query.unflatten(2, (attn.heads, -1))
202 | key = key.unflatten(2, (attn.heads, -1))
203 | value = value.unflatten(2, (attn.heads, -1))
204 |
205 | query = attn.norm_q(query)
206 | key = attn.norm_k(key)
207 |
208 | txt_length = encoder_hidden_states.shape[1]
209 |
210 | query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], image_rotary_emb), query[:, -txt_length:]], dim=1)
211 | key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], image_rotary_emb), key[:, -txt_length:]], dim=1)
212 |
213 | hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
214 | hidden_states = hidden_states.flatten(-2)
215 |
216 | hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
217 |
218 | return hidden_states, encoder_hidden_states
219 |
220 |
221 | class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
222 | def __init__(self, embedding_dim, pooled_projection_dim):
223 | super().__init__()
224 |
225 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
226 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
227 | self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
228 | self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
229 |
230 | def forward(self, timestep, guidance, pooled_projection):
231 | timesteps_proj = self.time_proj(timestep)
232 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
233 |
234 | guidance_proj = self.time_proj(guidance)
235 | guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
236 |
237 | time_guidance_emb = timesteps_emb + guidance_emb
238 |
239 | pooled_projections = self.text_embedder(pooled_projection)
240 | conditioning = time_guidance_emb + pooled_projections
241 |
242 | return conditioning
243 |
244 |
245 | class CombinedTimestepTextProjEmbeddings(nn.Module):
246 | def __init__(self, embedding_dim, pooled_projection_dim):
247 | super().__init__()
248 |
249 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
250 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
251 | self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
252 |
253 | def forward(self, timestep, pooled_projection):
254 | timesteps_proj = self.time_proj(timestep)
255 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
256 |
257 | pooled_projections = self.text_embedder(pooled_projection)
258 |
259 | conditioning = timesteps_emb + pooled_projections
260 |
261 | return conditioning
262 |
263 |
264 | class HunyuanVideoAdaNorm(nn.Module):
265 | def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
266 | super().__init__()
267 |
268 | out_features = out_features or 2 * in_features
269 | self.linear = nn.Linear(in_features, out_features)
270 | self.nonlinearity = nn.SiLU()
271 |
272 | def forward(
273 | self, temb: torch.Tensor
274 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
275 | temb = self.linear(self.nonlinearity(temb))
276 | gate_msa, gate_mlp = temb.chunk(2, dim=-1)
277 | gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
278 | return gate_msa, gate_mlp
279 |
280 |
281 | class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
282 | def __init__(
283 | self,
284 | num_attention_heads: int,
285 | attention_head_dim: int,
286 | mlp_width_ratio: str = 4.0,
287 | mlp_drop_rate: float = 0.0,
288 | attention_bias: bool = True,
289 | ) -> None:
290 | super().__init__()
291 |
292 | hidden_size = num_attention_heads * attention_head_dim
293 |
294 | self.norm1 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
295 | self.attn = Attention(
296 | query_dim=hidden_size,
297 | cross_attention_dim=None,
298 | heads=num_attention_heads,
299 | dim_head=attention_head_dim,
300 | bias=attention_bias,
301 | )
302 |
303 | self.norm2 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
304 | self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
305 |
306 | self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
307 |
308 | def forward(
309 | self,
310 | hidden_states: torch.Tensor,
311 | temb: torch.Tensor,
312 | attention_mask: Optional[torch.Tensor] = None,
313 | ) -> torch.Tensor:
314 | norm_hidden_states = self.norm1(hidden_states)
315 |
316 | attn_output = self.attn(
317 | hidden_states=norm_hidden_states,
318 | encoder_hidden_states=None,
319 | attention_mask=attention_mask,
320 | )
321 |
322 | gate_msa, gate_mlp = self.norm_out(temb)
323 | hidden_states = hidden_states + attn_output * gate_msa
324 |
325 | ff_output = self.ff(self.norm2(hidden_states))
326 | hidden_states = hidden_states + ff_output * gate_mlp
327 |
328 | return hidden_states
329 |
330 |
331 | class HunyuanVideoIndividualTokenRefiner(nn.Module):
332 | def __init__(
333 | self,
334 | num_attention_heads: int,
335 | attention_head_dim: int,
336 | num_layers: int,
337 | mlp_width_ratio: float = 4.0,
338 | mlp_drop_rate: float = 0.0,
339 | attention_bias: bool = True,
340 | ) -> None:
341 | super().__init__()
342 |
343 | self.refiner_blocks = nn.ModuleList(
344 | [
345 | HunyuanVideoIndividualTokenRefinerBlock(
346 | num_attention_heads=num_attention_heads,
347 | attention_head_dim=attention_head_dim,
348 | mlp_width_ratio=mlp_width_ratio,
349 | mlp_drop_rate=mlp_drop_rate,
350 | attention_bias=attention_bias,
351 | )
352 | for _ in range(num_layers)
353 | ]
354 | )
355 |
356 | def forward(
357 | self,
358 | hidden_states: torch.Tensor,
359 | temb: torch.Tensor,
360 | attention_mask: Optional[torch.Tensor] = None,
361 | ) -> None:
362 | self_attn_mask = None
363 | if attention_mask is not None:
364 | batch_size = attention_mask.shape[0]
365 | seq_len = attention_mask.shape[1]
366 | attention_mask = attention_mask.to(hidden_states.device).bool()
367 | self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
368 | self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
369 | self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
370 | self_attn_mask[:, :, :, 0] = True
371 |
372 | for block in self.refiner_blocks:
373 | hidden_states = block(hidden_states, temb, self_attn_mask)
374 |
375 | return hidden_states
376 |
377 |
378 | class HunyuanVideoTokenRefiner(nn.Module):
379 | def __init__(
380 | self,
381 | in_channels: int,
382 | num_attention_heads: int,
383 | attention_head_dim: int,
384 | num_layers: int,
385 | mlp_ratio: float = 4.0,
386 | mlp_drop_rate: float = 0.0,
387 | attention_bias: bool = True,
388 | ) -> None:
389 | super().__init__()
390 |
391 | hidden_size = num_attention_heads * attention_head_dim
392 |
393 | self.time_text_embed = CombinedTimestepTextProjEmbeddings(
394 | embedding_dim=hidden_size, pooled_projection_dim=in_channels
395 | )
396 | self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
397 | self.token_refiner = HunyuanVideoIndividualTokenRefiner(
398 | num_attention_heads=num_attention_heads,
399 | attention_head_dim=attention_head_dim,
400 | num_layers=num_layers,
401 | mlp_width_ratio=mlp_ratio,
402 | mlp_drop_rate=mlp_drop_rate,
403 | attention_bias=attention_bias,
404 | )
405 |
406 | def forward(
407 | self,
408 | hidden_states: torch.Tensor,
409 | timestep: torch.LongTensor,
410 | attention_mask: Optional[torch.LongTensor] = None,
411 | ) -> torch.Tensor:
412 | if attention_mask is None:
413 | pooled_projections = hidden_states.mean(dim=1)
414 | else:
415 | original_dtype = hidden_states.dtype
416 | mask_float = attention_mask.float().unsqueeze(-1)
417 | pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
418 | pooled_projections = pooled_projections.to(original_dtype)
419 |
420 | temb = self.time_text_embed(timestep, pooled_projections)
421 | hidden_states = self.proj_in(hidden_states)
422 | hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
423 |
424 | return hidden_states
425 |
426 |
427 | class HunyuanVideoRotaryPosEmbed(nn.Module):
428 | def __init__(self, rope_dim, theta):
429 | super().__init__()
430 | self.DT, self.DY, self.DX = rope_dim
431 | self.theta = theta
432 |
433 | @torch.no_grad()
434 | def get_frequency(self, dim, pos):
435 | T, H, W = pos.shape
436 | freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device)[: (dim // 2)] / dim))
437 | freqs = torch.outer(freqs, pos.reshape(-1)).unflatten(-1, (T, H, W)).repeat_interleave(2, dim=0)
438 | return freqs.cos(), freqs.sin()
439 |
440 | @torch.no_grad()
441 | def forward_inner(self, frame_indices, height, width, device):
442 | GT, GY, GX = torch.meshgrid(
443 | frame_indices.to(device=device, dtype=torch.float32),
444 | torch.arange(0, height, device=device, dtype=torch.float32),
445 | torch.arange(0, width, device=device, dtype=torch.float32),
446 | indexing="ij"
447 | )
448 |
449 | FCT, FST = self.get_frequency(self.DT, GT)
450 | FCY, FSY = self.get_frequency(self.DY, GY)
451 | FCX, FSX = self.get_frequency(self.DX, GX)
452 |
453 | result = torch.cat([FCT, FCY, FCX, FST, FSY, FSX], dim=0)
454 |
455 | return result.to(device)
456 |
457 | @torch.no_grad()
458 | def forward(self, frame_indices, height, width, device):
459 | frame_indices = frame_indices.unbind(0)
460 | results = [self.forward_inner(f, height, width, device) for f in frame_indices]
461 | results = torch.stack(results, dim=0)
462 | return results
463 |
464 |
465 | class AdaLayerNormZero(nn.Module):
466 | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
467 | super().__init__()
468 | self.silu = nn.SiLU()
469 | self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias)
470 | if norm_type == "layer_norm":
471 | self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
472 | else:
473 | raise ValueError(f"unknown norm_type {norm_type}")
474 |
475 | def forward(
476 | self,
477 | x: torch.Tensor,
478 | emb: Optional[torch.Tensor] = None,
479 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
480 | emb = emb.unsqueeze(-2)
481 | emb = self.linear(self.silu(emb))
482 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
483 | x = self.norm(x) * (1 + scale_msa) + shift_msa
484 | return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
485 |
486 |
487 | class AdaLayerNormZeroSingle(nn.Module):
488 | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
489 | super().__init__()
490 |
491 | self.silu = nn.SiLU()
492 | self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
493 | if norm_type == "layer_norm":
494 | self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
495 | else:
496 | raise ValueError(f"unknown norm_type {norm_type}")
497 |
498 | def forward(
499 | self,
500 | x: torch.Tensor,
501 | emb: Optional[torch.Tensor] = None,
502 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
503 | emb = emb.unsqueeze(-2)
504 | emb = self.linear(self.silu(emb))
505 | shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1)
506 | x = self.norm(x) * (1 + scale_msa) + shift_msa
507 | return x, gate_msa
508 |
509 |
510 | class AdaLayerNormContinuous(nn.Module):
511 | def __init__(
512 | self,
513 | embedding_dim: int,
514 | conditioning_embedding_dim: int,
515 | elementwise_affine=True,
516 | eps=1e-5,
517 | bias=True,
518 | norm_type="layer_norm",
519 | ):
520 | super().__init__()
521 | self.silu = nn.SiLU()
522 | self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
523 | if norm_type == "layer_norm":
524 | self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
525 | else:
526 | raise ValueError(f"unknown norm_type {norm_type}")
527 |
528 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
529 | emb = emb.unsqueeze(-2)
530 | emb = self.linear(self.silu(emb))
531 | scale, shift = emb.chunk(2, dim=-1)
532 | x = self.norm(x) * (1 + scale) + shift
533 | return x
534 |
535 |
536 | class HunyuanVideoSingleTransformerBlock(nn.Module):
537 | def __init__(
538 | self,
539 | num_attention_heads: int,
540 | attention_head_dim: int,
541 | mlp_ratio: float = 4.0,
542 | qk_norm: str = "rms_norm",
543 | ) -> None:
544 | super().__init__()
545 |
546 | hidden_size = num_attention_heads * attention_head_dim
547 | mlp_dim = int(hidden_size * mlp_ratio)
548 |
549 | self.attn = Attention(
550 | query_dim=hidden_size,
551 | cross_attention_dim=None,
552 | dim_head=attention_head_dim,
553 | heads=num_attention_heads,
554 | out_dim=hidden_size,
555 | bias=True,
556 | processor=HunyuanAttnProcessorFlashAttnSingle(),
557 | qk_norm=qk_norm,
558 | eps=1e-6,
559 | pre_only=True,
560 | )
561 |
562 | self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
563 | self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
564 | self.act_mlp = nn.GELU(approximate="tanh")
565 | self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
566 |
567 | def forward(
568 | self,
569 | hidden_states: torch.Tensor,
570 | encoder_hidden_states: torch.Tensor,
571 | temb: torch.Tensor,
572 | attention_mask: Optional[torch.Tensor] = None,
573 | image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
574 | ) -> torch.Tensor:
575 | text_seq_length = encoder_hidden_states.shape[1]
576 | hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
577 |
578 | residual = hidden_states
579 |
580 | # 1. Input normalization
581 | norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
582 | mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
583 |
584 | norm_hidden_states, norm_encoder_hidden_states = (
585 | norm_hidden_states[:, :-text_seq_length, :],
586 | norm_hidden_states[:, -text_seq_length:, :],
587 | )
588 |
589 | # 2. Attention
590 | attn_output, context_attn_output = self.attn(
591 | hidden_states=norm_hidden_states,
592 | encoder_hidden_states=norm_encoder_hidden_states,
593 | attention_mask=attention_mask,
594 | image_rotary_emb=image_rotary_emb,
595 | )
596 | attn_output = torch.cat([attn_output, context_attn_output], dim=1)
597 |
598 | # 3. Modulation and residual connection
599 | hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
600 | hidden_states = gate * self.proj_out(hidden_states)
601 | hidden_states = hidden_states + residual
602 |
603 | hidden_states, encoder_hidden_states = (
604 | hidden_states[:, :-text_seq_length, :],
605 | hidden_states[:, -text_seq_length:, :],
606 | )
607 | return hidden_states, encoder_hidden_states
608 |
609 |
610 | class HunyuanVideoTransformerBlock(nn.Module):
611 | def __init__(
612 | self,
613 | num_attention_heads: int,
614 | attention_head_dim: int,
615 | mlp_ratio: float,
616 | qk_norm: str = "rms_norm",
617 | ) -> None:
618 | super().__init__()
619 |
620 | hidden_size = num_attention_heads * attention_head_dim
621 |
622 | self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
623 | self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
624 |
625 | self.attn = Attention(
626 | query_dim=hidden_size,
627 | cross_attention_dim=None,
628 | added_kv_proj_dim=hidden_size,
629 | dim_head=attention_head_dim,
630 | heads=num_attention_heads,
631 | out_dim=hidden_size,
632 | context_pre_only=False,
633 | bias=True,
634 | processor=HunyuanAttnProcessorFlashAttnDouble(),
635 | qk_norm=qk_norm,
636 | eps=1e-6,
637 | )
638 |
639 | self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
640 | self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
641 |
642 | self.norm2_context = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
643 | self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
644 |
645 | def forward(
646 | self,
647 | hidden_states: torch.Tensor,
648 | encoder_hidden_states: torch.Tensor,
649 | temb: torch.Tensor,
650 | attention_mask: Optional[torch.Tensor] = None,
651 | freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
652 | ) -> Tuple[torch.Tensor, torch.Tensor]:
653 | # 1. Input normalization
654 | norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
655 | norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb)
656 |
657 | # 2. Joint attention
658 | attn_output, context_attn_output = self.attn(
659 | hidden_states=norm_hidden_states,
660 | encoder_hidden_states=norm_encoder_hidden_states,
661 | attention_mask=attention_mask,
662 | image_rotary_emb=freqs_cis,
663 | )
664 |
665 | # 3. Modulation and residual connection
666 | hidden_states = hidden_states + attn_output * gate_msa
667 | encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa
668 |
669 | norm_hidden_states = self.norm2(hidden_states)
670 | norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
671 |
672 | norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
673 | norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
674 |
675 | # 4. Feed-forward
676 | ff_output = self.ff(norm_hidden_states)
677 | context_ff_output = self.ff_context(norm_encoder_hidden_states)
678 |
679 | hidden_states = hidden_states + gate_mlp * ff_output
680 | encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
681 |
682 | return hidden_states, encoder_hidden_states
683 |
684 |
685 | class ClipVisionProjection(nn.Module):
686 | def __init__(self, in_channels, out_channels):
687 | super().__init__()
688 | self.up = nn.Linear(in_channels, out_channels * 3)
689 | self.down = nn.Linear(out_channels * 3, out_channels)
690 |
691 | def forward(self, x):
692 | projected_x = self.down(nn.functional.silu(self.up(x)))
693 | return projected_x
694 |
695 |
696 | class HunyuanVideoPatchEmbed(nn.Module):
697 | def __init__(self, patch_size, in_chans, embed_dim):
698 | super().__init__()
699 | self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
700 |
701 |
702 | class HunyuanVideoPatchEmbedForCleanLatents(nn.Module):
703 | def __init__(self, inner_dim):
704 | super().__init__()
705 | self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
706 | self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
707 | self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
708 |
709 | @torch.no_grad()
710 | def initialize_weight_from_another_conv3d(self, another_layer):
711 | weight = another_layer.weight.detach().clone()
712 | bias = another_layer.bias.detach().clone()
713 |
714 | sd = {
715 | 'proj.weight': weight.clone(),
716 | 'proj.bias': bias.clone(),
717 | 'proj_2x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=2, hk=2, wk=2) / 8.0,
718 | 'proj_2x.bias': bias.clone(),
719 | 'proj_4x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=4, hk=4, wk=4) / 64.0,
720 | 'proj_4x.bias': bias.clone(),
721 | }
722 |
723 | sd = {k: v.clone() for k, v in sd.items()}
724 |
725 | self.load_state_dict(sd)
726 | return
727 |
728 |
729 | class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
730 | @register_to_config
731 | def __init__(
732 | self,
733 | in_channels: int = 16,
734 | out_channels: int = 16,
735 | num_attention_heads: int = 24,
736 | attention_head_dim: int = 128,
737 | num_layers: int = 20,
738 | num_single_layers: int = 40,
739 | num_refiner_layers: int = 2,
740 | mlp_ratio: float = 4.0,
741 | patch_size: int = 2,
742 | patch_size_t: int = 1,
743 | qk_norm: str = "rms_norm",
744 | guidance_embeds: bool = True,
745 | text_embed_dim: int = 4096,
746 | pooled_projection_dim: int = 768,
747 | rope_theta: float = 256.0,
748 | rope_axes_dim: Tuple[int] = (16, 56, 56),
749 | has_image_proj=False,
750 | image_proj_dim=1152,
751 | has_clean_x_embedder=False,
752 | ) -> None:
753 | super().__init__()
754 |
755 | inner_dim = num_attention_heads * attention_head_dim
756 | out_channels = out_channels or in_channels
757 |
758 | # 1. Latent and condition embedders
759 | self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
760 | self.context_embedder = HunyuanVideoTokenRefiner(
761 | text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
762 | )
763 | self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
764 |
765 | self.clean_x_embedder = None
766 | self.image_projection = None
767 |
768 | # 2. RoPE
769 | self.rope = HunyuanVideoRotaryPosEmbed(rope_axes_dim, rope_theta)
770 |
771 | # 3. Dual stream transformer blocks
772 | self.transformer_blocks = nn.ModuleList(
773 | [
774 | HunyuanVideoTransformerBlock(
775 | num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
776 | )
777 | for _ in range(num_layers)
778 | ]
779 | )
780 |
781 | # 4. Single stream transformer blocks
782 | self.single_transformer_blocks = nn.ModuleList(
783 | [
784 | HunyuanVideoSingleTransformerBlock(
785 | num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
786 | )
787 | for _ in range(num_single_layers)
788 | ]
789 | )
790 |
791 | # 5. Output projection
792 | self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
793 | self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
794 |
795 | self.inner_dim = inner_dim
796 | self.use_gradient_checkpointing = False
797 | self.enable_teacache = False
798 |
799 | if has_image_proj:
800 | self.install_image_projection(image_proj_dim)
801 |
802 | if has_clean_x_embedder:
803 | self.install_clean_x_embedder()
804 |
805 | self.high_quality_fp32_output_for_inference = False
806 |
807 | def install_image_projection(self, in_channels):
808 | self.image_projection = ClipVisionProjection(in_channels=in_channels, out_channels=self.inner_dim)
809 | self.config['has_image_proj'] = True
810 | self.config['image_proj_dim'] = in_channels
811 |
812 | def install_clean_x_embedder(self):
813 | self.clean_x_embedder = HunyuanVideoPatchEmbedForCleanLatents(self.inner_dim)
814 | self.config['has_clean_x_embedder'] = True
815 |
816 | def enable_gradient_checkpointing(self):
817 | self.use_gradient_checkpointing = True
818 | print('self.use_gradient_checkpointing = True')
819 |
820 | def disable_gradient_checkpointing(self):
821 | self.use_gradient_checkpointing = False
822 | print('self.use_gradient_checkpointing = False')
823 |
824 | def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15):
825 | self.enable_teacache = enable_teacache
826 | self.cnt = 0
827 | self.num_steps = num_steps
828 | self.rel_l1_thresh = rel_l1_thresh # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup
829 | self.accumulated_rel_l1_distance = 0
830 | self.previous_modulated_input = None
831 | self.previous_residual = None
832 | self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02])
833 |
834 | def gradient_checkpointing_method(self, block, *args):
835 | if self.use_gradient_checkpointing:
836 | result = torch.utils.checkpoint.checkpoint(block, *args, use_reentrant=False)
837 | else:
838 | result = block(*args)
839 | return result
840 |
841 | def process_input_hidden_states(
842 | self,
843 | latents, latent_indices=None,
844 | clean_latents=None, clean_latent_indices=None,
845 | clean_latents_2x=None, clean_latent_2x_indices=None,
846 | clean_latents_4x=None, clean_latent_4x_indices=None
847 | ):
848 | hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents)
849 | B, C, T, H, W = hidden_states.shape
850 |
851 | if latent_indices is None:
852 | latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1)
853 |
854 | hidden_states = hidden_states.flatten(2).transpose(1, 2)
855 |
856 | rope_freqs = self.rope(frame_indices=latent_indices, height=H, width=W, device=hidden_states.device)
857 | rope_freqs = rope_freqs.flatten(2).transpose(1, 2)
858 |
859 | if clean_latents is not None and clean_latent_indices is not None:
860 | clean_latents = clean_latents.to(hidden_states)
861 | clean_latents = self.gradient_checkpointing_method(self.clean_x_embedder.proj, clean_latents)
862 | clean_latents = clean_latents.flatten(2).transpose(1, 2)
863 |
864 | clean_latent_rope_freqs = self.rope(frame_indices=clean_latent_indices, height=H, width=W, device=clean_latents.device)
865 | clean_latent_rope_freqs = clean_latent_rope_freqs.flatten(2).transpose(1, 2)
866 |
867 | hidden_states = torch.cat([clean_latents, hidden_states], dim=1)
868 | rope_freqs = torch.cat([clean_latent_rope_freqs, rope_freqs], dim=1)
869 |
870 | if clean_latents_2x is not None and clean_latent_2x_indices is not None:
871 | clean_latents_2x = clean_latents_2x.to(hidden_states)
872 | clean_latents_2x = pad_for_3d_conv(clean_latents_2x, (2, 4, 4))
873 | clean_latents_2x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_2x, clean_latents_2x)
874 | clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2)
875 |
876 | clean_latent_2x_rope_freqs = self.rope(frame_indices=clean_latent_2x_indices, height=H, width=W, device=clean_latents_2x.device)
877 | clean_latent_2x_rope_freqs = pad_for_3d_conv(clean_latent_2x_rope_freqs, (2, 2, 2))
878 | clean_latent_2x_rope_freqs = center_down_sample_3d(clean_latent_2x_rope_freqs, (2, 2, 2))
879 | clean_latent_2x_rope_freqs = clean_latent_2x_rope_freqs.flatten(2).transpose(1, 2)
880 |
881 | hidden_states = torch.cat([clean_latents_2x, hidden_states], dim=1)
882 | rope_freqs = torch.cat([clean_latent_2x_rope_freqs, rope_freqs], dim=1)
883 |
884 | if clean_latents_4x is not None and clean_latent_4x_indices is not None:
885 | clean_latents_4x = clean_latents_4x.to(hidden_states)
886 | clean_latents_4x = pad_for_3d_conv(clean_latents_4x, (4, 8, 8))
887 | clean_latents_4x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_4x, clean_latents_4x)
888 | clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2)
889 |
890 | clean_latent_4x_rope_freqs = self.rope(frame_indices=clean_latent_4x_indices, height=H, width=W, device=clean_latents_4x.device)
891 | clean_latent_4x_rope_freqs = pad_for_3d_conv(clean_latent_4x_rope_freqs, (4, 4, 4))
892 | clean_latent_4x_rope_freqs = center_down_sample_3d(clean_latent_4x_rope_freqs, (4, 4, 4))
893 | clean_latent_4x_rope_freqs = clean_latent_4x_rope_freqs.flatten(2).transpose(1, 2)
894 |
895 | hidden_states = torch.cat([clean_latents_4x, hidden_states], dim=1)
896 | rope_freqs = torch.cat([clean_latent_4x_rope_freqs, rope_freqs], dim=1)
897 |
898 | return hidden_states, rope_freqs
899 |
900 | def forward(
901 | self,
902 | hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance,
903 | latent_indices=None,
904 | clean_latents=None, clean_latent_indices=None,
905 | clean_latents_2x=None, clean_latent_2x_indices=None,
906 | clean_latents_4x=None, clean_latent_4x_indices=None,
907 | image_embeddings=None,
908 | attention_kwargs=None, return_dict=True
909 | ):
910 |
911 | if attention_kwargs is None:
912 | attention_kwargs = {}
913 |
914 | batch_size, num_channels, num_frames, height, width = hidden_states.shape
915 | p, p_t = self.config['patch_size'], self.config['patch_size_t']
916 | post_patch_num_frames = num_frames // p_t
917 | post_patch_height = height // p
918 | post_patch_width = width // p
919 | original_context_length = post_patch_num_frames * post_patch_height * post_patch_width
920 |
921 | hidden_states, rope_freqs = self.process_input_hidden_states(hidden_states, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices, clean_latents_4x, clean_latent_4x_indices)
922 |
923 | temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections)
924 | encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask)
925 |
926 | if self.image_projection is not None:
927 | assert image_embeddings is not None, 'You must use image embeddings!'
928 | extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings)
929 | extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device)
930 |
931 | # must cat before (not after) encoder_hidden_states, due to attn masking
932 | encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1)
933 | encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1)
934 |
935 | if batch_size == 1:
936 | # When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want
937 | # If they are not same, then their impls are wrong. Ours are always the correct one.
938 | text_len = encoder_attention_mask.sum().item()
939 | encoder_hidden_states = encoder_hidden_states[:, :text_len]
940 | attention_mask = None, None, None, None
941 | else:
942 | img_seq_len = hidden_states.shape[1]
943 | txt_seq_len = encoder_hidden_states.shape[1]
944 |
945 | cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len)
946 | cu_seqlens_kv = cu_seqlens_q
947 | max_seqlen_q = img_seq_len + txt_seq_len
948 | max_seqlen_kv = max_seqlen_q
949 |
950 | attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv
951 |
952 | if self.enable_teacache:
953 | modulated_inp = self.transformer_blocks[0].norm1(hidden_states, emb=temb)[0]
954 |
955 | if self.cnt == 0 or self.cnt == self.num_steps-1:
956 | should_calc = True
957 | self.accumulated_rel_l1_distance = 0
958 | else:
959 | curr_rel_l1 = ((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()
960 | self.accumulated_rel_l1_distance += self.teacache_rescale_func(curr_rel_l1)
961 | should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh
962 |
963 | if should_calc:
964 | self.accumulated_rel_l1_distance = 0
965 |
966 | self.previous_modulated_input = modulated_inp
967 | self.cnt += 1
968 |
969 | if self.cnt == self.num_steps:
970 | self.cnt = 0
971 |
972 | if not should_calc:
973 | hidden_states = hidden_states + self.previous_residual
974 | else:
975 | ori_hidden_states = hidden_states.clone()
976 |
977 | for block_id, block in enumerate(self.transformer_blocks):
978 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
979 | block,
980 | hidden_states,
981 | encoder_hidden_states,
982 | temb,
983 | attention_mask,
984 | rope_freqs
985 | )
986 |
987 | for block_id, block in enumerate(self.single_transformer_blocks):
988 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
989 | block,
990 | hidden_states,
991 | encoder_hidden_states,
992 | temb,
993 | attention_mask,
994 | rope_freqs
995 | )
996 |
997 | self.previous_residual = hidden_states - ori_hidden_states
998 | else:
999 | for block_id, block in enumerate(self.transformer_blocks):
1000 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
1001 | block,
1002 | hidden_states,
1003 | encoder_hidden_states,
1004 | temb,
1005 | attention_mask,
1006 | rope_freqs
1007 | )
1008 |
1009 | for block_id, block in enumerate(self.single_transformer_blocks):
1010 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
1011 | block,
1012 | hidden_states,
1013 | encoder_hidden_states,
1014 | temb,
1015 | attention_mask,
1016 | rope_freqs
1017 | )
1018 |
1019 | hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb)
1020 |
1021 | hidden_states = hidden_states[:, -original_context_length:, :]
1022 |
1023 | if self.high_quality_fp32_output_for_inference:
1024 | hidden_states = hidden_states.to(dtype=torch.float32)
1025 | if self.proj_out.weight.dtype != torch.float32:
1026 | self.proj_out.to(dtype=torch.float32)
1027 |
1028 | hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states)
1029 |
1030 | hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)',
1031 | t=post_patch_num_frames, h=post_patch_height, w=post_patch_width,
1032 | pt=p_t, ph=p, pw=p)
1033 |
1034 | if return_dict:
1035 | return Transformer2DModelOutput(sample=hidden_states)
1036 |
1037 | return hidden_states,
1038 |
--------------------------------------------------------------------------------