Upload 27 files

.gitignore

@@ -0,0 +1,182 @@
+hf_download/
+outputs/
+repo/
+loras/
+queue.json
+settings.json
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+temp/

README.md

@@ -0,0 +1,87 @@
+# FramePack Studio
+
+FramePack Studio is an enhanced version of the FramePack demo script, designed to create intricate video scenes with improved prompt adherence. This is very much a work in progress, expect some bugs and broken features. 
+![screencapture-127-0-0-1-7860-2025-05-04-20_13_58](https://github.com/user-attachments/assets/8fcb90af-8c3f-47ca-8f23-61d9b59438ae)
+
+
+## Current Features
+
+- **F1 and Original FramePack Models**: Run both in a single queue
+- **Timestamped Prompts**: Define different prompts for specific time segments in your video
+- **Prompt Blending**: Define the blending time between timestamped prompts
+- **Basic LoRA Support**: Works with most (all?) hunyuan LoRAs
+- **Queue System**: Process multiple generation jobs without blocking the interface
+- **Metadata Saving/Import**: Prompt and seed are encoded into the output PNG, all other generation metadata is saved in a JSON file
+- **I2V and T2V**: Works with or without an input image to allow for more flexibility when working with standard LoRAs
+- **Latent Image Options**: When using T2V you can generate based on a black, white, green screen or pure noise image
+
+
+## Fresh Installation
+
+### Prerequisites
+
+- Python 3.10+
+- CUDA-compatible GPU with at least 8GB VRAM (16GB+ recommended)
+
+### Setup
+
+Install via the Pinokio community script "FP-Studio" or:
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/colinurbs/FramePack-Studio.git
+   cd FramePack-Studio
+   ```
+
+2. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+## Usage
+
+Run the studio interface:
+
+```bash
+python studio.py
+```
+
+Additional command line options:
+- `--share`: Create a public Gradio link to share your interface
+- `--server`: Specify the server address (default: 0.0.0.0)
+- `--port`: Specify a custom port
+- `--inbrowser`: Automatically open the interface in your browser
+
+## LoRAs
+
+Add LoRAs to the /loras/ folder at the root of the installation. Select the LoRAs you wish to load and set the weights for each generation.
+
+NOTE: slow lora loading is a known issue
+
+## Working with Timestamped Prompts
+
+You can create videos with changing prompts over time using the following syntax:
+
+```
+[0s: A serene forest with sunlight filtering through the trees ]
+[5s: A deer appears in the clearing ]
+[10s: The deer drinks from a small stream ]
+```
+
+Each timestamp defines when that prompt should start influencing the generation. The system will (hopefully) smoothly transition between prompts for a cohesive video.
+
+## Credits
+Many thanks to [Lvmin Zhang](https://github.com/lllyasviel) for the absolutely amazing work on the original [FramePack](https://github.com/lllyasviel/FramePack) code!
+
+Thanks to [Rickard Edén](https://github.com/neph1) for the LoRA code and their general contributions to this growing FramePack scene!
+
+Thanks to everyone who has joined the Discord, reported a bug, sumbitted a PR or helped with testing!
+
+
+
+    @article{zhang2025framepack,
+        title={Packing Input Frame Contexts in Next-Frame Prediction Models for Video Generation},
+        author={Lvmin Zhang and Maneesh Agrawala},
+        journal={Arxiv},
+        year={2025}
+    }

diffusers_helper/bucket_tools.py

@@ -0,0 +1,94 @@
+bucket_options = {
+    128: [
+        (96, 160),
+        (112, 144),
+        (128, 128),
+        (144, 112),
+        (160, 96),
+    ],
+    256: [
+        (192, 320),
+        (224, 288),
+        (256, 256),
+        (288, 224),
+        (320, 192),
+    ],
+    384: [
+        (256, 512),
+        (320, 448),
+        (384, 384),
+        (448, 320),
+        (512, 256),
+    ],
+    512: [
+        (352, 704),
+        (384, 640),
+        (448, 576),
+        (512, 512),
+        (576, 448),
+        (640, 384),
+        (704, 352),
+    ],
+    640: [
+        (416, 960),
+        (448, 864),
+        (480, 832),
+        (512, 768),
+        (544, 704),
+        (576, 672),
+        (608, 640),
+        (640, 640),
+        (640, 608),
+        (672, 576),
+        (704, 544),
+        (768, 512),
+        (832, 480),
+        (864, 448),
+        (960, 416),
+    ],
+    768: [
+        (512, 1024),
+        (576, 896),
+        (640, 832),
+        (704, 768),
+        (768, 768),
+        (768, 704),
+        (832, 640),
+        (896, 576),
+        (1024, 512),
+    ],
+}
+
+
+def find_nearest_bucket(h, w, resolution=640):
+    # Use the provided resolution or find the closest available bucket size
+    # print(f"find_nearest_bucket called with h={h}, w={w}, resolution={resolution}")
+    
+    if resolution not in bucket_options:
+        # Find the closest available resolution
+        available_resolutions = list(bucket_options.keys())
+        closest_resolution = min(available_resolutions, key=lambda x: abs(x - resolution))
+        # print(f"Resolution {resolution} not found in bucket options, using closest available: {closest_resolution}")
+        resolution = closest_resolution
+    # else:
+        # print(f"Resolution {resolution} found in bucket options")
+    
+    # Calculate the aspect ratio of the input image
+    input_aspect_ratio = w / h if h > 0 else 1.0
+    # print(f"Input aspect ratio: {input_aspect_ratio:.4f}")
+    
+    min_diff = float('inf')
+    best_bucket = None
+    
+    # Find the bucket size with the closest aspect ratio to the input image
+    for (bucket_h, bucket_w) in bucket_options[resolution]:
+        bucket_aspect_ratio = bucket_w / bucket_h if bucket_h > 0 else 1.0
+        # Calculate the difference in aspect ratios
+        diff = abs(bucket_aspect_ratio - input_aspect_ratio)
+        if diff < min_diff:
+            min_diff = diff
+            best_bucket = (bucket_h, bucket_w)
+        # print(f"  Checking bucket ({bucket_h}, {bucket_w}), aspect ratio={bucket_aspect_ratio:.4f}, diff={diff:.4f}, current best={best_bucket}")
+    
+    # print(f"Using resolution {resolution}, selected bucket: {best_bucket}")
+    return best_bucket

diffusers_helper/clip_vision.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+
+def hf_clip_vision_encode(image, feature_extractor, image_encoder):
+    assert isinstance(image, np.ndarray)
+    assert image.ndim == 3 and image.shape[2] == 3
+    assert image.dtype == np.uint8
+
+    preprocessed = feature_extractor.preprocess(images=image, return_tensors="pt").to(device=image_encoder.device, dtype=image_encoder.dtype)
+    image_encoder_output = image_encoder(**preprocessed)
+
+    return image_encoder_output

diffusers_helper/dit_common.py

@@ -0,0 +1,53 @@
+import torch
+import accelerate.accelerator
+
+from diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm, AdaLayerNormContinuous
+
+
+accelerate.accelerator.convert_outputs_to_fp32 = lambda x: x
+
+
+def LayerNorm_forward(self, x):
+    return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps).to(x)
+
+
+LayerNorm.forward = LayerNorm_forward
+torch.nn.LayerNorm.forward = LayerNorm_forward
+
+
+def FP32LayerNorm_forward(self, x):
+    origin_dtype = x.dtype
+    return torch.nn.functional.layer_norm(
+        x.float(),
+        self.normalized_shape,
+        self.weight.float() if self.weight is not None else None,
+        self.bias.float() if self.bias is not None else None,
+        self.eps,
+    ).to(origin_dtype)
+
+
+FP32LayerNorm.forward = FP32LayerNorm_forward
+
+
+def RMSNorm_forward(self, hidden_states):
+    input_dtype = hidden_states.dtype
+    variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+
+    if self.weight is None:
+        return hidden_states.to(input_dtype)
+
+    return hidden_states.to(input_dtype) * self.weight.to(input_dtype)
+
+
+RMSNorm.forward = RMSNorm_forward
+
+
+def AdaLayerNormContinuous_forward(self, x, conditioning_embedding):
+    emb = self.linear(self.silu(conditioning_embedding))
+    scale, shift = emb.chunk(2, dim=1)
+    x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+    return x
+
+
+AdaLayerNormContinuous.forward = AdaLayerNormContinuous_forward

diffusers_helper/gradio/progress_bar.py

@@ -0,0 +1,86 @@
+progress_html = '''
+<div class="loader-container">
+  <div class="loader"></div>
+  <div class="progress-container">
+    <progress value="*number*" max="100"></progress>
+  </div>
+  <span>*text*</span>
+</div>
+'''
+
+css = '''
+.loader-container {
+  display: flex; /* Use flex to align items horizontally */
+  align-items: center; /* Center items vertically within the container */
+  white-space: nowrap; /* Prevent line breaks within the container */
+}
+
+.loader {
+  border: 8px solid #f3f3f3; /* Light grey */
+  border-top: 8px solid #3498db; /* Blue */
+  border-radius: 50%;
+  width: 30px;
+  height: 30px;
+  animation: spin 2s linear infinite;
+}
+
+@keyframes spin {
+  0% { transform: rotate(0deg); }
+  100% { transform: rotate(360deg); }
+}
+
+/* Style the progress bar */
+progress {
+  appearance: none; /* Remove default styling */
+  height: 20px; /* Set the height of the progress bar */
+  border-radius: 5px; /* Round the corners of the progress bar */
+  background-color: #f3f3f3; /* Light grey background */
+  width: 100%;
+  vertical-align: middle !important;
+}
+
+/* Style the progress bar container */
+.progress-container {
+  margin-left: 20px;
+  margin-right: 20px;
+  flex-grow: 1; /* Allow the progress container to take up remaining space */
+}
+
+/* Set the color of the progress bar fill */
+progress::-webkit-progress-value {
+  background-color: #3498db; /* Blue color for the fill */
+}
+
+progress::-moz-progress-bar {
+  background-color: #3498db; /* Blue color for the fill in Firefox */
+}
+
+/* Style the text on the progress bar */
+progress::after {
+  content: attr(value '%'); /* Display the progress value followed by '%' */
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  color: white; /* Set text color */
+  font-size: 14px; /* Set font size */
+}
+
+/* Style other texts */
+.loader-container > span {
+  margin-left: 5px; /* Add spacing between the progress bar and the text */
+}
+
+.no-generating-animation > .generating {
+  display: none !important;
+}
+
+'''
+
+
+def make_progress_bar_html(number, text):
+    return progress_html.replace('*number*', str(number)).replace('*text*', text)
+
+
+def make_progress_bar_css():
+    return css

diffusers_helper/hf_login.py

@@ -0,0 +1,21 @@
+import os
+
+
+def login(token):
+    from huggingface_hub import login
+    import time
+
+    while True:
+        try:
+            login(token)
+            print('HF login ok.')
+            break
+        except Exception as e:
+            print(f'HF login failed: {e}. Retrying')
+            time.sleep(0.5)
+
+
+hf_token = os.environ.get('HF_TOKEN', None)
+
+if hf_token is not None:
+    login(hf_token)

diffusers_helper/hunyuan.py

@@ -0,0 +1,111 @@
+import torch
+
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
+from diffusers_helper.utils import crop_or_pad_yield_mask
+
+
+@torch.no_grad()
+def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, max_length=256):
+    assert isinstance(prompt, str)
+
+    prompt = [prompt]
+
+    # LLAMA
+
+    prompt_llama = [DEFAULT_PROMPT_TEMPLATE["template"].format(p) for p in prompt]
+    crop_start = DEFAULT_PROMPT_TEMPLATE["crop_start"]
+
+    llama_inputs = tokenizer(
+        prompt_llama,
+        padding="max_length",
+        max_length=max_length + crop_start,
+        truncation=True,
+        return_tensors="pt",
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_attention_mask=True,
+    )
+
+    llama_input_ids = llama_inputs.input_ids.to(text_encoder.device)
+    llama_attention_mask = llama_inputs.attention_mask.to(text_encoder.device)
+    llama_attention_length = int(llama_attention_mask.sum())
+
+    llama_outputs = text_encoder(
+        input_ids=llama_input_ids,
+        attention_mask=llama_attention_mask,
+        output_hidden_states=True,
+    )
+
+    llama_vec = llama_outputs.hidden_states[-3][:, crop_start:llama_attention_length]
+    # llama_vec_remaining = llama_outputs.hidden_states[-3][:, llama_attention_length:]
+    llama_attention_mask = llama_attention_mask[:, crop_start:llama_attention_length]
+
+    assert torch.all(llama_attention_mask.bool())
+
+    # CLIP
+
+    clip_l_input_ids = tokenizer_2(
+        prompt,
+        padding="max_length",
+        max_length=77,
+        truncation=True,
+        return_overflowing_tokens=False,
+        return_length=False,
+        return_tensors="pt",
+    ).input_ids
+    clip_l_pooler = text_encoder_2(clip_l_input_ids.to(text_encoder_2.device), output_hidden_states=False).pooler_output
+
+    return llama_vec, clip_l_pooler
+
+
+@torch.no_grad()
+def vae_decode_fake(latents):
+    latent_rgb_factors = [
+        [-0.0395, -0.0331, 0.0445],
+        [0.0696, 0.0795, 0.0518],
+        [0.0135, -0.0945, -0.0282],
+        [0.0108, -0.0250, -0.0765],
+        [-0.0209, 0.0032, 0.0224],
+        [-0.0804, -0.0254, -0.0639],
+        [-0.0991, 0.0271, -0.0669],
+        [-0.0646, -0.0422, -0.0400],
+        [-0.0696, -0.0595, -0.0894],
+        [-0.0799, -0.0208, -0.0375],
+        [0.1166, 0.1627, 0.0962],
+        [0.1165, 0.0432, 0.0407],
+        [-0.2315, -0.1920, -0.1355],
+        [-0.0270, 0.0401, -0.0821],
+        [-0.0616, -0.0997, -0.0727],
+        [0.0249, -0.0469, -0.1703]
+    ]  # From comfyui
+
+    latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761]
+
+    weight = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1)[:, :, None, None, None]
+    bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype)
+
+    images = torch.nn.functional.conv3d(latents, weight, bias=bias, stride=1, padding=0, dilation=1, groups=1)
+    images = images.clamp(0.0, 1.0)
+
+    return images
+
+
+@torch.no_grad()
+def vae_decode(latents, vae, image_mode=False):
+    latents = latents / vae.config.scaling_factor
+
+    if not image_mode:
+        image = vae.decode(latents.to(device=vae.device, dtype=vae.dtype)).sample
+    else:
+        latents = latents.to(device=vae.device, dtype=vae.dtype).unbind(2)
+        image = [vae.decode(l.unsqueeze(2)).sample for l in latents]
+        image = torch.cat(image, dim=2)
+
+    return image
+
+
+@torch.no_grad()
+def vae_encode(image, vae):
+    latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()
+    latents = latents * vae.config.scaling_factor
+    return latents

diffusers_helper/k_diffusion/uni_pc_fm.py

@@ -0,0 +1,141 @@
+# Better Flow Matching UniPC by Lvmin Zhang
+# (c) 2025
+# CC BY-SA 4.0
+# Attribution-ShareAlike 4.0 International Licence
+
+
+import torch
+
+from tqdm.auto import trange
+
+
+def expand_dims(v, dims):
+    return v[(...,) + (None,) * (dims - 1)]
+
+
+class FlowMatchUniPC:
+    def __init__(self, model, extra_args, variant='bh1'):
+        self.model = model
+        self.variant = variant
+        self.extra_args = extra_args
+
+    def model_fn(self, x, t):
+        return self.model(x, t, **self.extra_args)
+
+    def update_fn(self, x, model_prev_list, t_prev_list, t, order):
+        assert order <= len(model_prev_list)
+        dims = x.dim()
+
+        t_prev_0 = t_prev_list[-1]
+        lambda_prev_0 = - torch.log(t_prev_0)
+        lambda_t = - torch.log(t)
+        model_prev_0 = model_prev_list[-1]
+
+        h = lambda_t - lambda_prev_0
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            t_prev_i = t_prev_list[-(i + 1)]
+            model_prev_i = model_prev_list[-(i + 1)]
+            lambda_prev_i = - torch.log(t_prev_i)
+            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
+            rks.append(rk)
+            D1s.append((model_prev_i - model_prev_0) / rk)
+
+        rks.append(1.)
+        rks = torch.tensor(rks, device=x.device)
+
+        R = []
+        b = []
+
+        hh = -h[0]
+        h_phi_1 = torch.expm1(hh)
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.variant == 'bh1':
+            B_h = hh
+        elif self.variant == 'bh2':
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError('Bad variant!')
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= (i + 1)
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=x.device)
+
+        use_predictor = len(D1s) > 0
+
+        if use_predictor:
+            D1s = torch.stack(D1s, dim=1)
+            if order == 2:
+                rhos_p = torch.tensor([0.5], device=b.device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
+        else:
+            D1s = None
+            rhos_p = None
+
+        if order == 1:
+            rhos_c = torch.tensor([0.5], device=b.device)
+        else:
+            rhos_c = torch.linalg.solve(R, b)
+
+        x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
+
+        if use_predictor:
+            pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
+        else:
+            pred_res = 0
+
+        x_t = x_t_ - expand_dims(B_h, dims) * pred_res
+        model_t = self.model_fn(x_t, t)
+
+        if D1s is not None:
+            corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
+        else:
+            corr_res = 0
+
+        D1_t = (model_t - model_prev_0)
+        x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
+
+        return x_t, model_t
+
+    def sample(self, x, sigmas, callback=None, disable_pbar=False):
+        order = min(3, len(sigmas) - 2)
+        model_prev_list, t_prev_list = [], []
+        for i in trange(len(sigmas) - 1, disable=disable_pbar):
+            vec_t = sigmas[i].expand(x.shape[0])
+
+            if i == 0:
+                model_prev_list = [self.model_fn(x, vec_t)]
+                t_prev_list = [vec_t]
+            elif i < order:
+                init_order = i
+                x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
+                model_prev_list.append(model_x)
+                t_prev_list.append(vec_t)
+            else:
+                x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
+                model_prev_list.append(model_x)
+                t_prev_list.append(vec_t)
+
+            model_prev_list = model_prev_list[-order:]
+            t_prev_list = t_prev_list[-order:]
+
+            if callback is not None:
+                callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
+
+        return model_prev_list[-1]
+
+
+def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
+    assert variant in ['bh1', 'bh2']
+    return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)

diffusers_helper/k_diffusion/wrapper.py

@@ -0,0 +1,51 @@
+import torch
+
+
+def append_dims(x, target_dims):
+    return x[(...,) + (None,) * (target_dims - x.ndim)]
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
+    if guidance_rescale == 0:
+        return noise_cfg
+
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+def fm_wrapper(transformer, t_scale=1000.0):
+    def k_model(x, sigma, **extra_args):
+        dtype = extra_args['dtype']
+        cfg_scale = extra_args['cfg_scale']
+        cfg_rescale = extra_args['cfg_rescale']
+        concat_latent = extra_args['concat_latent']
+
+        original_dtype = x.dtype
+        sigma = sigma.float()
+
+        x = x.to(dtype)
+        timestep = (sigma * t_scale).to(dtype)
+
+        if concat_latent is None:
+            hidden_states = x
+        else:
+            hidden_states = torch.cat([x, concat_latent.to(x)], dim=1)
+
+        pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float()
+
+        if cfg_scale == 1.0:
+            pred_negative = torch.zeros_like(pred_positive)
+        else:
+            pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float()
+
+        pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative)
+        pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale)
+
+        x0 = x.float() - pred.float() * append_dims(sigma, x.ndim)
+
+        return x0.to(dtype=original_dtype)
+
+    return k_model

diffusers_helper/lora_utils.py

@@ -0,0 +1,131 @@
+from pathlib import Path, PurePath
+from typing import Dict, List, Optional, Union
+from diffusers.loaders.lora_pipeline import _fetch_state_dict
+from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers
+from diffusers.utils.peft_utils import set_weights_and_activate_adapters
+from diffusers.loaders.peft import _SET_ADAPTER_SCALE_FN_MAPPING
+import torch
+
+def load_lora(transformer, lora_path: Path, weight_name: Optional[str] = "pytorch_lora_weights.safetensors"):
+    """
+    Load LoRA weights into the transformer model.
+
+    Args:
+        transformer: The transformer model to which LoRA weights will be applied.
+        lora_path (Path): Path to the LoRA weights file.
+        weight_name (Optional[str]): Name of the weight to load.
+
+    """
+    
+    state_dict = _fetch_state_dict(
+        lora_path,
+        weight_name,
+        True,
+        True,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None)
+
+    state_dict = _convert_hunyuan_video_lora_to_diffusers(state_dict)
+    
+    # should weight_name even be Optional[str] or just str?
+    # For now, we assume it is never None
+    # The module name in the state_dict must not include a . in the name
+    # See https://github.com/pytorch/pytorch/pull/6639/files#diff-4be56271f7bfe650e3521c81fd363da58f109cd23ee80d243156d2d6ccda6263R133-R134
+    adapter_name = PurePath(str(weight_name).replace('_DOT_', '.')).stem.replace('.', '_DOT_')
+    if '_DOT_' in adapter_name:
+        print(
+            f"LoRA file '{weight_name}' contains a '.' in the name. " +
+            'This may cause issues. Consider renaming the file.' +
+            f" Using '{adapter_name}' as the adapter name to be safe."
+        )
+    
+    # Check if adapter already exists and delete it if it does
+    if hasattr(transformer, 'peft_config') and adapter_name in transformer.peft_config:
+        print(f"Adapter '{adapter_name}' already exists. Removing it before loading again.")
+        # Use delete_adapters (plural) instead of delete_adapter
+        transformer.delete_adapters([adapter_name])
+    
+    # Load the adapter with the original name
+    transformer.load_lora_adapter(state_dict, network_alphas=None, adapter_name=adapter_name)
+    print(f"LoRA weights '{adapter_name}' loaded successfully.")
+    
+    return transformer
+
+def unload_all_loras(transformer):
+    """
+    Completely unload all LoRA adapters from the transformer model.
+    """
+    if hasattr(transformer, 'peft_config') and transformer.peft_config:
+        # Get all adapter names
+        adapter_names = list(transformer.peft_config.keys())
+        
+        if adapter_names:
+            print(f"Removing all LoRA adapters: {', '.join(adapter_names)}")
+            # Delete all adapters
+            transformer.delete_adapters(adapter_names)
+            
+            # Force cleanup of any remaining adapter references
+            if hasattr(transformer, 'active_adapter'):
+                transformer.active_adapter = None
+                
+            # Clear any cached states
+            for module in transformer.modules():
+                if hasattr(module, 'lora_A'):
+                    if isinstance(module.lora_A, dict):
+                        module.lora_A.clear()
+                if hasattr(module, 'lora_B'):
+                    if isinstance(module.lora_B, dict):
+                        module.lora_B.clear()
+                if hasattr(module, 'scaling'):
+                    if isinstance(module.scaling, dict):
+                        module.scaling.clear()
+            
+            print("All LoRA adapters have been completely removed.")
+        else:
+            print("No LoRA adapters found to remove.")
+    else:
+        print("Model doesn't have any LoRA adapters or peft_config.")
+    
+    # Force garbage collection
+    import gc
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    return transformer
+
+    
+# TODO(neph1): remove when HunyuanVideoTransformer3DModelPacked is in _SET_ADAPTER_SCALE_FN_MAPPING
+def set_adapters(
+        transformer,
+        adapter_names: Union[List[str], str],
+        weights: Optional[Union[float, Dict, List[float], List[Dict], List[None]]] = None,
+    ):
+
+    adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names
+
+    # Expand weights into a list, one entry per adapter
+    # examples for e.g. 2 adapters:  [{...}, 7] -> [7,7] ; None -> [None, None]
+    if not isinstance(weights, list):
+        weights = [weights] * len(adapter_names)
+
+    if len(adapter_names) != len(weights):
+        raise ValueError(
+            f"Length of adapter names {len(adapter_names)} is not equal to the length of their weights {len(weights)}."
+        )
+
+    # Set None values to default of 1.0
+    # e.g. [{...}, 7] -> [{...}, 7] ; [None, None] -> [1.0, 1.0]
+    weights = [w if w is not None else 1.0 for w in weights]
+
+    # e.g. [{...}, 7] -> [{expanded dict...}, 7]
+    scale_expansion_fn = _SET_ADAPTER_SCALE_FN_MAPPING["HunyuanVideoTransformer3DModel"]
+    weights = scale_expansion_fn(transformer, weights)
+
+    set_weights_and_activate_adapters(transformer, adapter_names, weights)

diffusers_helper/memory.py

@@ -0,0 +1,134 @@
+# By lllyasviel
+
+
+import torch
+
+
+cpu = torch.device('cpu')
+gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(module: torch.nn.Module, **kwargs):
+        original_class = module.__class__
+        module.__dict__['forge_backup_original_class'] = original_class
+
+        def hacked_get_attr(self, name: str):
+            if '_parameters' in self.__dict__:
+                _parameters = self.__dict__['_parameters']
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if '_buffers' in self.__dict__:
+                _buffers = self.__dict__['_buffers']
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
+            '__getattr__': hacked_get_attr,
+        })
+
+        return
+
+    @staticmethod
+    def _uninstall_module(module: torch.nn.Module):
+        if 'forge_backup_original_class' in module.__dict__:
+            module.__class__ = module.__dict__.pop('forge_backup_original_class')
+        return
+
+    @staticmethod
+    def install_model(model: torch.nn.Module, **kwargs):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(model: torch.nn.Module):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+    if hasattr(model, 'scale_shift_table'):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, 'weight'):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(device=None):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats['active_bytes.all.current']
+    bytes_reserved = memory_stats['reserved_bytes.all.current']
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024 ** 3)
+
+
+def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(*args):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f'Unloaded {m.__class__.__name__} as complete.')
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(model, target_device, unload=True):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
+
+    gpu_complete_modules.append(model)
+    return

diffusers_helper/models/hunyuan_video_packed.py

@@ -0,0 +1,1032 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import einops
+import torch.nn as nn
+import numpy as np
+
+from diffusers.loaders import FromOriginalModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.utils import logging
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps, PixArtAlphaTextProjection
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers_helper.dit_common import LayerNorm
+from diffusers_helper.utils import zero_module
+
+
+enabled_backends = []
+
+if torch.backends.cuda.flash_sdp_enabled():
+    enabled_backends.append("flash")
+if torch.backends.cuda.math_sdp_enabled():
+    enabled_backends.append("math")
+if torch.backends.cuda.mem_efficient_sdp_enabled():
+    enabled_backends.append("mem_efficient")
+if torch.backends.cuda.cudnn_sdp_enabled():
+    enabled_backends.append("cudnn")
+
+print("Currently enabled native sdp backends:", enabled_backends)
+
+try:
+    # raise NotImplementedError
+    from xformers.ops import memory_efficient_attention as xformers_attn_func
+    print('Xformers is installed!')
+except:
+    print('Xformers is not installed!')
+    xformers_attn_func = None
+
+try:
+    # raise NotImplementedError
+    from flash_attn import flash_attn_varlen_func, flash_attn_func
+    print('Flash Attn is installed!')
+except:
+    print('Flash Attn is not installed!')
+    flash_attn_varlen_func = None
+    flash_attn_func = None
+
+try:
+    # raise NotImplementedError
+    from sageattention import sageattn_varlen, sageattn
+    print('Sage Attn is installed!')
+except:
+    print('Sage Attn is not installed!')
+    sageattn_varlen = None
+    sageattn = None
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def pad_for_3d_conv(x, kernel_size):
+    b, c, t, h, w = x.shape
+    pt, ph, pw = kernel_size
+    pad_t = (pt - (t % pt)) % pt
+    pad_h = (ph - (h % ph)) % ph
+    pad_w = (pw - (w % pw)) % pw
+    return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode='replicate')
+
+
+def center_down_sample_3d(x, kernel_size):
+    # pt, ph, pw = kernel_size
+    # cp = (pt * ph * pw) // 2
+    # 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)
+    # xc = xp[cp]
+    # return xc
+    return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size)
+
+
+def get_cu_seqlens(text_mask, img_len):
+    batch_size = text_mask.shape[0]
+    text_len = text_mask.sum(dim=1)
+    max_len = text_mask.shape[1] + img_len
+
+    cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda")
+
+    for i in range(batch_size):
+        s = text_len[i] + img_len
+        s1 = i * max_len + s
+        s2 = (i + 1) * max_len
+        cu_seqlens[2 * i + 1] = s1
+        cu_seqlens[2 * i + 2] = s2
+
+    return cu_seqlens
+
+
+def apply_rotary_emb_transposed(x, freqs_cis):
+    cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
+    x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1)
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+    out = x.float() * cos + x_rotated.float() * sin
+    out = out.to(x)
+    return out
+
+
+def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv):
+    if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None:
+        if sageattn is not None:
+            x = sageattn(q, k, v, tensor_layout='NHD')
+            return x
+
+        if flash_attn_func is not None:
+            x = flash_attn_func(q, k, v)
+            return x
+
+        if xformers_attn_func is not None:
+            x = xformers_attn_func(q, k, v)
+            return x
+
+        x = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
+        return x
+
+    batch_size = q.shape[0]
+    q = q.view(q.shape[0] * q.shape[1], *q.shape[2:])
+    k = k.view(k.shape[0] * k.shape[1], *k.shape[2:])
+    v = v.view(v.shape[0] * v.shape[1], *v.shape[2:])
+    if sageattn_varlen is not None:
+        x = sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
+    elif flash_attn_varlen_func is not None:
+        x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
+    else:
+        raise NotImplementedError('No Attn Installed!')
+    x = x.view(batch_size, max_seqlen_q, *x.shape[2:])
+    return x
+
+
+class HunyuanAttnProcessorFlashAttnDouble:
+    def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
+        cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(2, (attn.heads, -1))
+        key = key.unflatten(2, (attn.heads, -1))
+        value = value.unflatten(2, (attn.heads, -1))
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        query = apply_rotary_emb_transposed(query, image_rotary_emb)
+        key = apply_rotary_emb_transposed(key, image_rotary_emb)
+
+        encoder_query = attn.add_q_proj(encoder_hidden_states)
+        encoder_key = attn.add_k_proj(encoder_hidden_states)
+        encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+        encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
+        encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
+        encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
+
+        encoder_query = attn.norm_added_q(encoder_query)
+        encoder_key = attn.norm_added_k(encoder_key)
+
+        query = torch.cat([query, encoder_query], dim=1)
+        key = torch.cat([key, encoder_key], dim=1)
+        value = torch.cat([value, encoder_value], dim=1)
+
+        hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
+        hidden_states = hidden_states.flatten(-2)
+
+        txt_length = encoder_hidden_states.shape[1]
+        hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
+
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanAttnProcessorFlashAttnSingle:
+    def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
+        cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
+
+        hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(2, (attn.heads, -1))
+        key = key.unflatten(2, (attn.heads, -1))
+        value = value.unflatten(2, (attn.heads, -1))
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        txt_length = encoder_hidden_states.shape[1]
+
+        query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], image_rotary_emb), query[:, -txt_length:]], dim=1)
+        key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], image_rotary_emb), key[:, -txt_length:]], dim=1)
+
+        hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
+        hidden_states = hidden_states.flatten(-2)
+
+        hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
+
+        return hidden_states, encoder_hidden_states
+
+
+class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim, pooled_projection_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
+
+    def forward(self, timestep, guidance, pooled_projection):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
+
+        guidance_proj = self.time_proj(guidance)
+        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
+
+        time_guidance_emb = timesteps_emb + guidance_emb
+
+        pooled_projections = self.text_embedder(pooled_projection)
+        conditioning = time_guidance_emb + pooled_projections
+
+        return conditioning
+
+
+class CombinedTimestepTextProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim, pooled_projection_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
+
+    def forward(self, timestep, pooled_projection):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
+
+        pooled_projections = self.text_embedder(pooled_projection)
+
+        conditioning = timesteps_emb + pooled_projections
+
+        return conditioning
+
+
+class HunyuanVideoAdaNorm(nn.Module):
+    def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
+        super().__init__()
+
+        out_features = out_features or 2 * in_features
+        self.linear = nn.Linear(in_features, out_features)
+        self.nonlinearity = nn.SiLU()
+
+    def forward(
+        self, temb: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        temb = self.linear(self.nonlinearity(temb))
+        gate_msa, gate_mlp = temb.chunk(2, dim=-1)
+        gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
+        return gate_msa, gate_mlp
+
+
+class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_width_ratio: str = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+        )
+
+        self.norm2 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+        self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
+
+        self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=attention_mask,
+        )
+
+        gate_msa, gate_mlp = self.norm_out(temb)
+        hidden_states = hidden_states + attn_output * gate_msa
+
+        ff_output = self.ff(self.norm2(hidden_states))
+        hidden_states = hidden_states + ff_output * gate_mlp
+
+        return hidden_states
+
+
+class HunyuanVideoIndividualTokenRefiner(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.refiner_blocks = nn.ModuleList(
+            [
+                HunyuanVideoIndividualTokenRefinerBlock(
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_width_ratio=mlp_width_ratio,
+                    mlp_drop_rate=mlp_drop_rate,
+                    attention_bias=attention_bias,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> None:
+        self_attn_mask = None
+        if attention_mask is not None:
+            batch_size = attention_mask.shape[0]
+            seq_len = attention_mask.shape[1]
+            attention_mask = attention_mask.to(hidden_states.device).bool()
+            self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
+            self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
+            self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
+            self_attn_mask[:, :, :, 0] = True
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, temb, self_attn_mask)
+
+        return hidden_states
+
+
+class HunyuanVideoTokenRefiner(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+            embedding_dim=hidden_size, pooled_projection_dim=in_channels
+        )
+        self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
+        self.token_refiner = HunyuanVideoIndividualTokenRefiner(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
+            mlp_width_ratio=mlp_ratio,
+            mlp_drop_rate=mlp_drop_rate,
+            attention_bias=attention_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            pooled_projections = hidden_states.mean(dim=1)
+        else:
+            original_dtype = hidden_states.dtype
+            mask_float = attention_mask.float().unsqueeze(-1)
+            pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            pooled_projections = pooled_projections.to(original_dtype)
+
+        temb = self.time_text_embed(timestep, pooled_projections)
+        hidden_states = self.proj_in(hidden_states)
+        hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
+
+        return hidden_states
+
+
+class HunyuanVideoRotaryPosEmbed(nn.Module):
+    def __init__(self, rope_dim, theta):
+        super().__init__()
+        self.DT, self.DY, self.DX = rope_dim
+        self.theta = theta
+
+    @torch.no_grad()
+    def get_frequency(self, dim, pos):
+        T, H, W = pos.shape
+        freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device)[: (dim // 2)] / dim))
+        freqs = torch.outer(freqs, pos.reshape(-1)).unflatten(-1, (T, H, W)).repeat_interleave(2, dim=0)
+        return freqs.cos(), freqs.sin()
+
+    @torch.no_grad()
+    def forward_inner(self, frame_indices, height, width, device):
+        GT, GY, GX = torch.meshgrid(
+            frame_indices.to(device=device, dtype=torch.float32),
+            torch.arange(0, height, device=device, dtype=torch.float32),
+            torch.arange(0, width, device=device, dtype=torch.float32),
+            indexing="ij"
+        )
+
+        FCT, FST = self.get_frequency(self.DT, GT)
+        FCY, FSY = self.get_frequency(self.DY, GY)
+        FCX, FSX = self.get_frequency(self.DX, GX)
+
+        result = torch.cat([FCT, FCY, FCX, FST, FSY, FSX], dim=0)
+
+        return result.to(device)
+
+    @torch.no_grad()
+    def forward(self, frame_indices, height, width, device):
+        frame_indices = frame_indices.unbind(0)
+        results = [self.forward_inner(f, height, width, device) for f in frame_indices]
+        results = torch.stack(results, dim=0)
+        return results
+
+
+class AdaLayerNormZero(nn.Module):
+    def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias)
+        if norm_type == "layer_norm":
+            self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb = emb.unsqueeze(-2)
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class AdaLayerNormZeroSingle(nn.Module):
+    def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
+        if norm_type == "layer_norm":
+            self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb = emb.unsqueeze(-2)
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa
+        return x, gate_msa
+
+
+class AdaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        if norm_type == "layer_norm":
+            self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        emb = emb.unsqueeze(-2)
+        emb = self.linear(self.silu(emb))
+        scale, shift = emb.chunk(2, dim=-1)
+        x = self.norm(x) * (1 + scale) + shift
+        return x
+
+
+class HunyuanVideoSingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        qk_norm: str = "rms_norm",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+        mlp_dim = int(hidden_size * mlp_ratio)
+
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=hidden_size,
+            bias=True,
+            processor=HunyuanAttnProcessorFlashAttnSingle(),
+            qk_norm=qk_norm,
+            eps=1e-6,
+            pre_only=True,
+        )
+
+        self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
+        self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+
+        residual = hidden_states
+
+        # 1. Input normalization
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+        norm_hidden_states, norm_encoder_hidden_states = (
+            norm_hidden_states[:, :-text_seq_length, :],
+            norm_hidden_states[:, -text_seq_length:, :],
+        )
+
+        # 2. Attention
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            attention_mask=attention_mask,
+            image_rotary_emb=image_rotary_emb,
+        )
+        attn_output = torch.cat([attn_output, context_attn_output], dim=1)
+
+        # 3. Modulation and residual connection
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = hidden_states + residual
+
+        hidden_states, encoder_hidden_states = (
+            hidden_states[:, :-text_seq_length, :],
+            hidden_states[:, -text_seq_length:, :],
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanVideoTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float,
+        qk_norm: str = "rms_norm",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+        self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+
+        self.attn = Attention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            added_kv_proj_dim=hidden_size,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=hidden_size,
+            context_pre_only=False,
+            bias=True,
+            processor=HunyuanAttnProcessorFlashAttnDouble(),
+            qk_norm=qk_norm,
+            eps=1e-6,
+        )
+
+        self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
+
+        self.norm2_context = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Input normalization
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb)
+
+        # 2. Joint attention
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            attention_mask=attention_mask,
+            image_rotary_emb=freqs_cis,
+        )
+
+        # 3. Modulation and residual connection
+        hidden_states = hidden_states + attn_output * gate_msa
+        encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
+
+        # 4. Feed-forward
+        ff_output = self.ff(norm_hidden_states)
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+
+        hidden_states = hidden_states + gate_mlp * ff_output
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
+
+        return hidden_states, encoder_hidden_states
+
+
+class ClipVisionProjection(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.up = nn.Linear(in_channels, out_channels * 3)
+        self.down = nn.Linear(out_channels * 3, out_channels)
+
+    def forward(self, x):
+        projected_x = self.down(nn.functional.silu(self.up(x)))
+        return projected_x
+
+
+class HunyuanVideoPatchEmbed(nn.Module):
+    def __init__(self, patch_size, in_chans, embed_dim):
+        super().__init__()
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+
+class HunyuanVideoPatchEmbedForCleanLatents(nn.Module):
+    def __init__(self, inner_dim):
+        super().__init__()
+        self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+        self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
+        self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
+
+    @torch.no_grad()
+    def initialize_weight_from_another_conv3d(self, another_layer):
+        weight = another_layer.weight.detach().clone()
+        bias = another_layer.bias.detach().clone()
+
+        sd = {
+            'proj.weight': weight.clone(),
+            'proj.bias': bias.clone(),
+            '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,
+            'proj_2x.bias': bias.clone(),
+            '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,
+            'proj_4x.bias': bias.clone(),
+        }
+
+        sd = {k: v.clone() for k, v in sd.items()}
+
+        self.load_state_dict(sd)
+        return
+
+
+class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 16,
+        out_channels: int = 16,
+        num_attention_heads: int = 24,
+        attention_head_dim: int = 128,
+        num_layers: int = 20,
+        num_single_layers: int = 40,
+        num_refiner_layers: int = 2,
+        mlp_ratio: float = 4.0,
+        patch_size: int = 2,
+        patch_size_t: int = 1,
+        qk_norm: str = "rms_norm",
+        guidance_embeds: bool = True,
+        text_embed_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        rope_theta: float = 256.0,
+        rope_axes_dim: Tuple[int] = (16, 56, 56),
+        has_image_proj=False,
+        image_proj_dim=1152,
+        has_clean_x_embedder=False,
+    ) -> None:
+        super().__init__()
+
+        inner_dim = num_attention_heads * attention_head_dim
+        out_channels = out_channels or in_channels
+
+        # 1. Latent and condition embedders
+        self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
+        self.context_embedder = HunyuanVideoTokenRefiner(
+            text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
+        )
+        self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
+
+        self.clean_x_embedder = None
+        self.image_projection = None
+
+        # 2. RoPE
+        self.rope = HunyuanVideoRotaryPosEmbed(rope_axes_dim, rope_theta)
+
+        # 3. Dual stream transformer blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                HunyuanVideoTransformerBlock(
+                    num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Single stream transformer blocks
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                HunyuanVideoSingleTransformerBlock(
+                    num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        # 5. Output projection
+        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
+
+        self.inner_dim = inner_dim
+        self.use_gradient_checkpointing = False
+        self.enable_teacache = False
+
+        if has_image_proj:
+            self.install_image_projection(image_proj_dim)
+
+        if has_clean_x_embedder:
+            self.install_clean_x_embedder()
+
+        self.high_quality_fp32_output_for_inference = False
+
+    def install_image_projection(self, in_channels):
+        self.image_projection = ClipVisionProjection(in_channels=in_channels, out_channels=self.inner_dim)
+        self.config['has_image_proj'] = True
+        self.config['image_proj_dim'] = in_channels
+
+    def install_clean_x_embedder(self):
+        self.clean_x_embedder = HunyuanVideoPatchEmbedForCleanLatents(self.inner_dim)
+        self.config['has_clean_x_embedder'] = True
+
+    def enable_gradient_checkpointing(self):
+        self.use_gradient_checkpointing = True
+        print('self.use_gradient_checkpointing = True')
+
+    def disable_gradient_checkpointing(self):
+        self.use_gradient_checkpointing = False
+        print('self.use_gradient_checkpointing = False')
+
+    def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15):
+        self.enable_teacache = enable_teacache
+        self.cnt = 0
+        self.num_steps = num_steps
+        self.rel_l1_thresh = rel_l1_thresh  # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup
+        self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = None
+        self.previous_residual = None
+        self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02])
+
+    def gradient_checkpointing_method(self, block, *args):
+        if self.use_gradient_checkpointing:
+            result = torch.utils.checkpoint.checkpoint(block, *args, use_reentrant=False)
+        else:
+            result = block(*args)
+        return result
+
+    def process_input_hidden_states(
+            self,
+            latents, latent_indices=None,
+            clean_latents=None, clean_latent_indices=None,
+            clean_latents_2x=None, clean_latent_2x_indices=None,
+            clean_latents_4x=None, clean_latent_4x_indices=None
+    ):
+        hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents)
+        B, C, T, H, W = hidden_states.shape
+
+        if latent_indices is None:
+            latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1)
+
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        rope_freqs = self.rope(frame_indices=latent_indices, height=H, width=W, device=hidden_states.device)
+        rope_freqs = rope_freqs.flatten(2).transpose(1, 2)
+
+        if clean_latents is not None and clean_latent_indices is not None:
+            clean_latents = clean_latents.to(hidden_states)
+            clean_latents = self.gradient_checkpointing_method(self.clean_x_embedder.proj, clean_latents)
+            clean_latents = clean_latents.flatten(2).transpose(1, 2)
+
+            clean_latent_rope_freqs = self.rope(frame_indices=clean_latent_indices, height=H, width=W, device=clean_latents.device)
+            clean_latent_rope_freqs = clean_latent_rope_freqs.flatten(2).transpose(1, 2)
+
+            hidden_states = torch.cat([clean_latents, hidden_states], dim=1)
+            rope_freqs = torch.cat([clean_latent_rope_freqs, rope_freqs], dim=1)
+
+        if clean_latents_2x is not None and clean_latent_2x_indices is not None:
+            clean_latents_2x = clean_latents_2x.to(hidden_states)
+            clean_latents_2x = pad_for_3d_conv(clean_latents_2x, (2, 4, 4))
+            clean_latents_2x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_2x, clean_latents_2x)
+            clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2)
+
+            clean_latent_2x_rope_freqs = self.rope(frame_indices=clean_latent_2x_indices, height=H, width=W, device=clean_latents_2x.device)
+            clean_latent_2x_rope_freqs = pad_for_3d_conv(clean_latent_2x_rope_freqs, (2, 2, 2))
+            clean_latent_2x_rope_freqs = center_down_sample_3d(clean_latent_2x_rope_freqs, (2, 2, 2))
+            clean_latent_2x_rope_freqs = clean_latent_2x_rope_freqs.flatten(2).transpose(1, 2)
+
+            hidden_states = torch.cat([clean_latents_2x, hidden_states], dim=1)
+            rope_freqs = torch.cat([clean_latent_2x_rope_freqs, rope_freqs], dim=1)
+
+        if clean_latents_4x is not None and clean_latent_4x_indices is not None:
+            clean_latents_4x = clean_latents_4x.to(hidden_states)
+            clean_latents_4x = pad_for_3d_conv(clean_latents_4x, (4, 8, 8))
+            clean_latents_4x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_4x, clean_latents_4x)
+            clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2)
+
+            clean_latent_4x_rope_freqs = self.rope(frame_indices=clean_latent_4x_indices, height=H, width=W, device=clean_latents_4x.device)
+            clean_latent_4x_rope_freqs = pad_for_3d_conv(clean_latent_4x_rope_freqs, (4, 4, 4))
+            clean_latent_4x_rope_freqs = center_down_sample_3d(clean_latent_4x_rope_freqs, (4, 4, 4))
+            clean_latent_4x_rope_freqs = clean_latent_4x_rope_freqs.flatten(2).transpose(1, 2)
+
+            hidden_states = torch.cat([clean_latents_4x, hidden_states], dim=1)
+            rope_freqs = torch.cat([clean_latent_4x_rope_freqs, rope_freqs], dim=1)
+
+        return hidden_states, rope_freqs
+
+    def forward(
+            self,
+            hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance,
+            latent_indices=None,
+            clean_latents=None, clean_latent_indices=None,
+            clean_latents_2x=None, clean_latent_2x_indices=None,
+            clean_latents_4x=None, clean_latent_4x_indices=None,
+            image_embeddings=None,
+            attention_kwargs=None, return_dict=True
+    ):
+
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p, p_t = self.config['patch_size'], self.config['patch_size_t']
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p
+        post_patch_width = width // p
+        original_context_length = post_patch_num_frames * post_patch_height * post_patch_width
+
+        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)
+
+        temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections)
+        encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask)
+
+        if self.image_projection is not None:
+            assert image_embeddings is not None, 'You must use image embeddings!'
+            extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings)
+            extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device)
+
+            # must cat before (not after) encoder_hidden_states, due to attn masking
+            encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1)
+            encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1)
+
+        with torch.no_grad():
+            if batch_size == 1:
+                # When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want
+                # If they are not same, then their impls are wrong. Ours are always the correct one.
+                text_len = encoder_attention_mask.sum().item()
+                encoder_hidden_states = encoder_hidden_states[:, :text_len]
+                attention_mask = None, None, None, None
+            else:
+                img_seq_len = hidden_states.shape[1]
+                txt_seq_len = encoder_hidden_states.shape[1]
+
+                cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len)
+                cu_seqlens_kv = cu_seqlens_q
+                max_seqlen_q = img_seq_len + txt_seq_len
+                max_seqlen_kv = max_seqlen_q
+
+                attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv
+
+        if self.enable_teacache:
+            modulated_inp = self.transformer_blocks[0].norm1(hidden_states, emb=temb)[0]
+
+            if self.cnt == 0 or self.cnt == self.num_steps-1:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+            else:
+                curr_rel_l1 = ((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()
+                self.accumulated_rel_l1_distance += self.teacache_rescale_func(curr_rel_l1)
+                should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh
+
+                if should_calc:
+                    self.accumulated_rel_l1_distance = 0
+
+            self.previous_modulated_input = modulated_inp
+            self.cnt += 1
+
+            if self.cnt == self.num_steps:
+                self.cnt = 0
+
+            if not should_calc:
+                hidden_states = hidden_states + self.previous_residual
+            else:
+                ori_hidden_states = hidden_states.clone()
+
+                for block_id, block in enumerate(self.transformer_blocks):
+                    hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
+                        block,
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        attention_mask,
+                        rope_freqs
+                    )
+
+                for block_id, block in enumerate(self.single_transformer_blocks):
+                    hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
+                        block,
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        attention_mask,
+                        rope_freqs
+                    )
+
+                self.previous_residual = hidden_states - ori_hidden_states
+        else:
+            for block_id, block in enumerate(self.transformer_blocks):
+                hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    attention_mask,
+                    rope_freqs
+                )
+
+            for block_id, block in enumerate(self.single_transformer_blocks):
+                hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    attention_mask,
+                    rope_freqs
+                )
+
+        hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb)
+
+        hidden_states = hidden_states[:, -original_context_length:, :]
+
+        if self.high_quality_fp32_output_for_inference:
+            hidden_states = hidden_states.to(dtype=torch.float32)
+            if self.proj_out.weight.dtype != torch.float32:
+                self.proj_out.to(dtype=torch.float32)
+
+        hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states)
+
+        hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)',
+                                         t=post_patch_num_frames, h=post_patch_height, w=post_patch_width,
+                                         pt=p_t, ph=p, pw=p)
+
+        if return_dict:
+            return Transformer2DModelOutput(sample=hidden_states)
+
+        return hidden_states,

diffusers_helper/pipelines/k_diffusion_hunyuan.py

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

diffusers_helper/thread_utils.py

@@ -0,0 +1,76 @@
+import time
+
+from threading import Thread, Lock
+
+
+class Listener:
+    task_queue = []
+    lock = Lock()
+    thread = None
+    
+    @classmethod
+    def _process_tasks(cls):
+        while True:
+            task = None
+            with cls.lock:
+                if cls.task_queue:
+                    task = cls.task_queue.pop(0)
+                    
+            if task is None:
+                time.sleep(0.001)
+                continue
+                
+            func, args, kwargs = task
+            try:
+                func(*args, **kwargs)
+            except Exception as e:
+                print(f"Error in listener thread: {e}")
+    
+    @classmethod
+    def add_task(cls, func, *args, **kwargs):
+        with cls.lock:
+            cls.task_queue.append((func, args, kwargs))
+
+        if cls.thread is None:
+            cls.thread = Thread(target=cls._process_tasks, daemon=True)
+            cls.thread.start()
+
+
+def async_run(func, *args, **kwargs):
+    Listener.add_task(func, *args, **kwargs)
+
+
+class FIFOQueue:
+    def __init__(self):
+        self.queue = []
+        self.lock = Lock()
+
+    def push(self, item):
+        with self.lock:
+            self.queue.append(item)
+
+    def pop(self):
+        with self.lock:
+            if self.queue:
+                return self.queue.pop(0)
+            return None
+
+    def top(self):
+        with self.lock:
+            if self.queue:
+                return self.queue[0]
+            return None
+
+    def next(self):
+        while True:
+            with self.lock:
+                if self.queue:
+                    return self.queue.pop(0)
+
+            time.sleep(0.001)
+
+
+class AsyncStream:
+    def __init__(self):
+        self.input_queue = FIFOQueue()
+        self.output_queue = FIFOQueue()

diffusers_helper/utils.py

@@ -0,0 +1,613 @@
+import os
+import cv2
+import json
+import random
+import glob
+import torch
+import einops
+import numpy as np
+import datetime
+import torchvision
+
+import safetensors.torch as sf
+from PIL import Image
+
+
+def min_resize(x, m):
+    if x.shape[0] < x.shape[1]:
+        s0 = m
+        s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
+    else:
+        s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
+        s1 = m
+    new_max = max(s1, s0)
+    raw_max = max(x.shape[0], x.shape[1])
+    if new_max < raw_max:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (s1, s0), interpolation=interpolation)
+    return y
+
+
+def d_resize(x, y):
+    H, W, C = y.shape
+    new_min = min(H, W)
+    raw_min = min(x.shape[0], x.shape[1])
+    if new_min < raw_min:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (W, H), interpolation=interpolation)
+    return y
+
+
+def resize_and_center_crop(image, target_width, target_height):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    original_width, original_height = pil_image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return np.array(cropped_image)
+
+
+def resize_and_center_crop_pytorch(image, target_width, target_height):
+    B, C, H, W = image.shape
+
+    if H == target_height and W == target_width:
+        return image
+
+    scale_factor = max(target_width / W, target_height / H)
+    resized_width = int(round(W * scale_factor))
+    resized_height = int(round(H * scale_factor))
+
+    resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
+
+    top = (resized_height - target_height) // 2
+    left = (resized_width - target_width) // 2
+    cropped = resized[:, :, top:top + target_height, left:left + target_width]
+
+    return cropped
+
+
+def resize_without_crop(image, target_width, target_height):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
+    return np.array(resized_image)
+
+
+def just_crop(image, w, h):
+    if h == image.shape[0] and w == image.shape[1]:
+        return image
+
+    original_height, original_width = image.shape[:2]
+    k = min(original_height / h, original_width / w)
+    new_width = int(round(w * k))
+    new_height = int(round(h * k))
+    x_start = (original_width - new_width) // 2
+    y_start = (original_height - new_height) // 2
+    cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
+    return cropped_image
+
+
+def write_to_json(data, file_path):
+    temp_file_path = file_path + ".tmp"
+    with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
+        json.dump(data, temp_file, indent=4)
+    os.replace(temp_file_path, file_path)
+    return
+
+
+def read_from_json(file_path):
+    with open(file_path, 'rt', encoding='utf-8') as file:
+        data = json.load(file)
+    return data
+
+
+def get_active_parameters(m):
+    return {k: v for k, v in m.named_parameters() if v.requires_grad}
+
+
+def cast_training_params(m, dtype=torch.float32):
+    result = {}
+    for n, param in m.named_parameters():
+        if param.requires_grad:
+            param.data = param.to(dtype)
+            result[n] = param
+    return result
+
+
+def separate_lora_AB(parameters, B_patterns=None):
+    parameters_normal = {}
+    parameters_B = {}
+
+    if B_patterns is None:
+        B_patterns = ['.lora_B.', '__zero__']
+
+    for k, v in parameters.items():
+        if any(B_pattern in k for B_pattern in B_patterns):
+            parameters_B[k] = v
+        else:
+            parameters_normal[k] = v
+
+    return parameters_normal, parameters_B
+
+
+def set_attr_recursive(obj, attr, value):
+    attrs = attr.split(".")
+    for name in attrs[:-1]:
+        obj = getattr(obj, name)
+    setattr(obj, attrs[-1], value)
+    return
+
+
+def print_tensor_list_size(tensors):
+    total_size = 0
+    total_elements = 0
+
+    if isinstance(tensors, dict):
+        tensors = tensors.values()
+
+    for tensor in tensors:
+        total_size += tensor.nelement() * tensor.element_size()
+        total_elements += tensor.nelement()
+
+    total_size_MB = total_size / (1024 ** 2)
+    total_elements_B = total_elements / 1e9
+
+    print(f"Total number of tensors: {len(tensors)}")
+    print(f"Total size of tensors: {total_size_MB:.2f} MB")
+    print(f"Total number of parameters: {total_elements_B:.3f} billion")
+    return
+
+
+@torch.no_grad()
+def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
+    batch_size = a.size(0)
+
+    if b is None:
+        b = torch.zeros_like(a)
+
+    if mask_a is None:
+        mask_a = torch.rand(batch_size) < probability_a
+
+    mask_a = mask_a.to(a.device)
+    mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
+    result = torch.where(mask_a, a, b)
+    return result
+
+
+@torch.no_grad()
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+@torch.no_grad()
+def supress_lower_channels(m, k, alpha=0.01):
+    data = m.weight.data.clone()
+
+    assert int(data.shape[1]) >= k
+
+    data[:, :k] = data[:, :k] * alpha
+    m.weight.data = data.contiguous().clone()
+    return m
+
+
+def freeze_module(m):
+    if not hasattr(m, '_forward_inside_frozen_module'):
+        m._forward_inside_frozen_module = m.forward
+    m.requires_grad_(False)
+    m.forward = torch.no_grad()(m.forward)
+    return m
+
+
+def get_latest_safetensors(folder_path):
+    safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
+
+    if not safetensors_files:
+        raise ValueError('No file to resume!')
+
+    latest_file = max(safetensors_files, key=os.path.getmtime)
+    latest_file = os.path.abspath(os.path.realpath(latest_file))
+    return latest_file
+
+
+def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
+    tags = tags_str.split(', ')
+    tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
+    prompt = ', '.join(tags)
+    return prompt
+
+
+def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
+    numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
+    edges = np.linspace(0, 1, n + 1)
+    points = np.random.uniform(edges[:-1], edges[1:])
+    numbers = inclusive + (exclusive - inclusive) * points
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def soft_append_bcthw(history, current, overlap=0):
+    if overlap <= 0:
+        return torch.cat([history, current], dim=2)
+
+    assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
+    assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
+    
+    weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
+    blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
+    output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
+
+    return output.to(history)
+
+
+def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
+    b, c, t, h, w = x.shape
+
+    per_row = b
+    for p in [6, 5, 4, 3, 2]:
+        if b % p == 0:
+            per_row = p
+            break
+
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
+    torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})
+    return x
+
+
+def save_bcthw_as_png(x, output_filename):
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def save_bchw_as_png(x, output_filename):
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, 'b c h w -> c h (b w)')
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def add_tensors_with_padding(tensor1, tensor2):
+    if tensor1.shape == tensor2.shape:
+        return tensor1 + tensor2
+
+    shape1 = tensor1.shape
+    shape2 = tensor2.shape
+
+    new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
+
+    padded_tensor1 = torch.zeros(new_shape)
+    padded_tensor2 = torch.zeros(new_shape)
+
+    padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
+    padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
+
+    result = padded_tensor1 + padded_tensor2
+    return result
+
+
+def print_free_mem():
+    torch.cuda.empty_cache()
+    free_mem, total_mem = torch.cuda.mem_get_info(0)
+    free_mem_mb = free_mem / (1024 ** 2)
+    total_mem_mb = total_mem / (1024 ** 2)
+    print(f"Free memory: {free_mem_mb:.2f} MB")
+    print(f"Total memory: {total_mem_mb:.2f} MB")
+    return
+
+
+def print_gpu_parameters(device, state_dict, log_count=1):
+    summary = {"device": device, "keys_count": len(state_dict)}
+
+    logged_params = {}
+    for i, (key, tensor) in enumerate(state_dict.items()):
+        if i >= log_count:
+            break
+        logged_params[key] = tensor.flatten()[:3].tolist()
+
+    summary["params"] = logged_params
+
+    print(str(summary))
+    return
+
+
+def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
+    from PIL import Image, ImageDraw, ImageFont
+
+    txt = Image.new("RGB", (width, height), color="white")
+    draw = ImageDraw.Draw(txt)
+    font = ImageFont.truetype(font_path, size=size)
+
+    if text == '':
+        return np.array(txt)
+
+    # Split text into lines that fit within the image width
+    lines = []
+    words = text.split()
+    current_line = words[0]
+
+    for word in words[1:]:
+        line_with_word = f"{current_line} {word}"
+        if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
+            current_line = line_with_word
+        else:
+            lines.append(current_line)
+            current_line = word
+
+    lines.append(current_line)
+
+    # Draw the text line by line
+    y = 0
+    line_height = draw.textbbox((0, 0), "A", font=font)[3]
+
+    for line in lines:
+        if y + line_height > height:
+            break  # stop drawing if the next line will be outside the image
+        draw.text((0, y), line, fill="black", font=font)
+        y += line_height
+
+    return np.array(txt)
+
+
+def blue_mark(x):
+    x = x.copy()
+    c = x[:, :, 2]
+    b = cv2.blur(c, (9, 9))
+    x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
+    return x
+
+
+def green_mark(x):
+    x = x.copy()
+    x[:, :, 2] = -1
+    x[:, :, 0] = -1
+    return x
+
+
+def frame_mark(x):
+    x = x.copy()
+    x[:64] = -1
+    x[-64:] = -1
+    x[:, :8] = 1
+    x[:, -8:] = 1
+    return x
+
+
+@torch.inference_mode()
+def pytorch2numpy(imgs):
+    results = []
+    for x in imgs:
+        y = x.movedim(0, -1)
+        y = y * 127.5 + 127.5
+        y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
+        results.append(y)
+    return results
+
+
+@torch.inference_mode()
+def numpy2pytorch(imgs):
+    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
+    h = h.movedim(-1, 1)
+    return h
+
+
+@torch.no_grad()
+def duplicate_prefix_to_suffix(x, count, zero_out=False):
+    if zero_out:
+        return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
+    else:
+        return torch.cat([x, x[:count]], dim=0)
+
+
+def weighted_mse(a, b, weight):
+    return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
+
+
+def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
+    x = (x - x_min) / (x_max - x_min)
+    x = max(0.0, min(x, 1.0))
+    x = x ** sigma
+    return y_min + x * (y_max - y_min)
+
+
+def expand_to_dims(x, target_dims):
+    return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
+
+
+def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
+    if tensor is None:
+        return None
+
+    first_dim = tensor.shape[0]
+
+    if first_dim == batch_size:
+        return tensor
+
+    if batch_size % first_dim != 0:
+        raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
+
+    repeat_times = batch_size // first_dim
+
+    return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
+
+
+def dim5(x):
+    return expand_to_dims(x, 5)
+
+
+def dim4(x):
+    return expand_to_dims(x, 4)
+
+
+def dim3(x):
+    return expand_to_dims(x, 3)
+
+
+def crop_or_pad_yield_mask(x, length):
+    B, F, C = x.shape
+    device = x.device
+    dtype = x.dtype
+
+    if F < length:
+        y = torch.zeros((B, length, C), dtype=dtype, device=device)
+        mask = torch.zeros((B, length), dtype=torch.bool, device=device)
+        y[:, :F, :] = x
+        mask[:, :F] = True
+        return y, mask
+
+    return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
+
+
+def extend_dim(x, dim, minimal_length, zero_pad=False):
+    original_length = int(x.shape[dim])
+
+    if original_length >= minimal_length:
+        return x
+
+    if zero_pad:
+        padding_shape = list(x.shape)
+        padding_shape[dim] = minimal_length - original_length
+        padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
+    else:
+        idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
+        last_element = x[idx]
+        padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
+
+    return torch.cat([x, padding], dim=dim)
+
+
+def lazy_positional_encoding(t, repeats=None):
+    if not isinstance(t, list):
+        t = [t]
+
+    from diffusers.models.embeddings import get_timestep_embedding
+
+    te = torch.tensor(t)
+    te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
+
+    if repeats is None:
+        return te
+
+    te = te[:, None, :].expand(-1, repeats, -1)
+
+    return te
+
+
+def state_dict_offset_merge(A, B, C=None):
+    result = {}
+    keys = A.keys()
+
+    for key in keys:
+        A_value = A[key]
+        B_value = B[key].to(A_value)
+
+        if C is None:
+            result[key] = A_value + B_value
+        else:
+            C_value = C[key].to(A_value)
+            result[key] = A_value + B_value - C_value
+
+    return result
+
+
+def state_dict_weighted_merge(state_dicts, weights):
+    if len(state_dicts) != len(weights):
+        raise ValueError("Number of state dictionaries must match number of weights")
+
+    if not state_dicts:
+        return {}
+
+    total_weight = sum(weights)
+
+    if total_weight == 0:
+        raise ValueError("Sum of weights cannot be zero")
+
+    normalized_weights = [w / total_weight for w in weights]
+
+    keys = state_dicts[0].keys()
+    result = {}
+
+    for key in keys:
+        result[key] = state_dicts[0][key] * normalized_weights[0]
+
+        for i in range(1, len(state_dicts)):
+            state_dict_value = state_dicts[i][key].to(result[key])
+            result[key] += state_dict_value * normalized_weights[i]
+
+    return result
+
+
+def group_files_by_folder(all_files):
+    grouped_files = {}
+
+    for file in all_files:
+        folder_name = os.path.basename(os.path.dirname(file))
+        if folder_name not in grouped_files:
+            grouped_files[folder_name] = []
+        grouped_files[folder_name].append(file)
+
+    list_of_lists = list(grouped_files.values())
+    return list_of_lists
+
+
+def generate_timestamp():
+    now = datetime.datetime.now()
+    timestamp = now.strftime('%y%m%d_%H%M%S')
+    milliseconds = f"{int(now.microsecond / 1000):03d}"
+    random_number = random.randint(0, 9999)
+    return f"{timestamp}_{milliseconds}_{random_number}"
+
+
+def write_PIL_image_with_png_info(image, metadata, path):
+    from PIL.PngImagePlugin import PngInfo
+
+    png_info = PngInfo()
+    for key, value in metadata.items():
+        png_info.add_text(key, value)
+
+    image.save(path, "PNG", pnginfo=png_info)
+    return image
+
+
+def torch_safe_save(content, path):
+    torch.save(content, path + '_tmp')
+    os.replace(path + '_tmp', path)
+    return path
+
+
+def move_optimizer_to_device(optimizer, device):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(device)

modules/__init__.py

@@ -0,0 +1,2 @@
+# modules/__init__.py
+

modules/generators/__init__.py

@@ -0,0 +1,23 @@
+from .original_generator import OriginalModelGenerator
+from .f1_generator import F1ModelGenerator
+
+def create_model_generator(model_type, **kwargs):
+    """
+    Create a model generator based on the model type.
+    
+    Args:
+        model_type: The type of model to create ("Original" or "F1")
+        **kwargs: Additional arguments to pass to the model generator constructor
+        
+    Returns:
+        A model generator instance
+        
+    Raises:
+        ValueError: If the model type is not supported
+    """
+    if model_type == "Original":
+        return OriginalModelGenerator(**kwargs)
+    elif model_type == "F1":
+        return F1ModelGenerator(**kwargs)
+    else:
+        raise ValueError(f"Unsupported model type: {model_type}")

modules/generators/base_generator.py

@@ -0,0 +1,228 @@
+import torch
+from abc import ABC, abstractmethod
+from diffusers_helper import lora_utils
+
+class BaseModelGenerator(ABC):
+    """
+    Base class for model generators.
+    This defines the common interface that all model generators must implement.
+    """
+    
+    def __init__(self, 
+                 text_encoder, 
+                 text_encoder_2, 
+                 tokenizer, 
+                 tokenizer_2, 
+                 vae, 
+                 image_encoder, 
+                 feature_extractor, 
+                 high_vram=False,
+                 prompt_embedding_cache=None,
+                 settings=None):
+        """
+        Initialize the base model generator.
+        
+        Args:
+            text_encoder: The text encoder model
+            text_encoder_2: The second text encoder model
+            tokenizer: The tokenizer for the first text encoder
+            tokenizer_2: The tokenizer for the second text encoder
+            vae: The VAE model
+            image_encoder: The image encoder model
+            feature_extractor: The feature extractor
+            high_vram: Whether high VRAM mode is enabled
+            prompt_embedding_cache: Cache for prompt embeddings
+            settings: Application settings
+        """
+        self.text_encoder = text_encoder
+        self.text_encoder_2 = text_encoder_2
+        self.tokenizer = tokenizer
+        self.tokenizer_2 = tokenizer_2
+        self.vae = vae
+        self.image_encoder = image_encoder
+        self.feature_extractor = feature_extractor
+        self.high_vram = high_vram
+        self.prompt_embedding_cache = prompt_embedding_cache or {}
+        self.settings = settings
+        self.transformer = None
+        self.gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.cpu = torch.device("cpu")
+    
+    @abstractmethod
+    def load_model(self):
+        """
+        Load the transformer model.
+        This method should be implemented by each specific model generator.
+        """
+        pass
+    
+    @abstractmethod
+    def get_model_name(self):
+        """
+        Get the name of the model.
+        This method should be implemented by each specific model generator.
+        """
+        pass
+    
+    def unload_loras(self):
+        """
+        Unload all LoRAs from the transformer model.
+        """
+        if self.transformer is not None:
+            print(f"Unloading all LoRAs from {self.get_model_name()} model")
+            self.transformer = lora_utils.unload_all_loras(self.transformer)
+            self.verify_lora_state("After unloading LoRAs")
+            import gc
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+    
+    def verify_lora_state(self, label=""):
+        """
+        Debug function to verify the state of LoRAs in the transformer model.
+        """
+        if self.transformer is None:
+            print(f"[{label}] Transformer is None, cannot verify LoRA state")
+            return
+            
+        has_loras = False
+        if hasattr(self.transformer, 'peft_config'):
+            adapter_names = list(self.transformer.peft_config.keys()) if self.transformer.peft_config else []
+            if adapter_names:
+                has_loras = True
+                print(f"[{label}] Transformer has LoRAs: {', '.join(adapter_names)}")
+            else:
+                print(f"[{label}] Transformer has no LoRAs in peft_config")
+        else:
+            print(f"[{label}] Transformer has no peft_config attribute")
+            
+        # Check for any LoRA modules
+        for name, module in self.transformer.named_modules():
+            if hasattr(module, 'lora_A') and module.lora_A:
+                has_loras = True
+                # print(f"[{label}] Found lora_A in module {name}")
+            if hasattr(module, 'lora_B') and module.lora_B:
+                has_loras = True
+                # print(f"[{label}] Found lora_B in module {name}")
+                
+        if not has_loras:
+            print(f"[{label}] No LoRA components found in transformer")
+    
+    def move_lora_adapters_to_device(self, target_device):
+        """
+        Move all LoRA adapters in the transformer model to the specified device.
+        This handles the PEFT implementation of LoRA.
+        """
+        if self.transformer is None:
+            return
+            
+        print(f"Moving all LoRA adapters to {target_device}")
+        
+        # First, find all modules with LoRA adapters
+        lora_modules = []
+        for name, module in self.transformer.named_modules():
+            if hasattr(module, 'active_adapter') and hasattr(module, 'lora_A') and hasattr(module, 'lora_B'):
+                lora_modules.append((name, module))
+        
+        # Now move all LoRA components to the target device
+        for name, module in lora_modules:
+            # Get the active adapter name
+            active_adapter = module.active_adapter
+            
+            # Move the LoRA layers to the target device
+            if active_adapter is not None:
+                if isinstance(module.lora_A, torch.nn.ModuleDict):
+                    # Handle ModuleDict case (PEFT implementation)
+                    for adapter_name in list(module.lora_A.keys()):
+                        # Move lora_A
+                        if adapter_name in module.lora_A:
+                            module.lora_A[adapter_name] = module.lora_A[adapter_name].to(target_device)
+                        
+                        # Move lora_B
+                        if adapter_name in module.lora_B:
+                            module.lora_B[adapter_name] = module.lora_B[adapter_name].to(target_device)
+                        
+                        # Move scaling
+                        if hasattr(module, 'scaling') and isinstance(module.scaling, dict) and adapter_name in module.scaling:
+                            if isinstance(module.scaling[adapter_name], torch.Tensor):
+                                module.scaling[adapter_name] = module.scaling[adapter_name].to(target_device)
+                else:
+                    # Handle direct attribute case
+                    if hasattr(module, 'lora_A') and module.lora_A is not None:
+                        module.lora_A = module.lora_A.to(target_device)
+                    if hasattr(module, 'lora_B') and module.lora_B is not None:
+                        module.lora_B = module.lora_B.to(target_device)
+                    if hasattr(module, 'scaling') and module.scaling is not None:
+                        if isinstance(module.scaling, torch.Tensor):
+                            module.scaling = module.scaling.to(target_device)
+        
+        print(f"Moved all LoRA adapters to {target_device}")
+    
+    def load_loras(self, selected_loras, lora_folder, lora_loaded_names, lora_values=None):
+        """
+        Load LoRAs into the transformer model.
+        
+        Args:
+            selected_loras: List of LoRA names to load
+            lora_folder: Folder containing the LoRA files
+            lora_loaded_names: List of loaded LoRA names
+            lora_values: Optional list of LoRA strength values
+        """
+        if self.transformer is None:
+            print("Cannot load LoRAs: Transformer model is not loaded")
+            return
+            
+        import os
+        
+        # Ensure all LoRAs are unloaded first
+        self.unload_loras()
+        
+        # Load each selected LoRA
+        for lora_name in selected_loras:
+            try:
+                idx = lora_loaded_names.index(lora_name)
+                lora_file = None
+                for ext in [".safetensors", ".pt"]:
+                    # Find any file that starts with the lora_name and ends with the extension
+                    matching_files = [f for f in os.listdir(lora_folder) 
+                                   if f.startswith(lora_name) and f.endswith(ext)]
+                    if matching_files:
+                        lora_file = matching_files[0]  # Use the first matching file
+                        break
+                        
+                if lora_file:
+                    print(f"Loading LoRA {lora_file} to {self.get_model_name()} model")
+                    self.transformer = lora_utils.load_lora(self.transformer, lora_folder, lora_file)
+                    
+                    # Set LoRA strength if provided
+                    if lora_values and idx < len(lora_values):
+                        lora_strength = float(lora_values[idx])
+                        print(f"Setting LoRA {lora_name} strength to {lora_strength}")
+                        
+                        # Set scaling for this LoRA by iterating through modules
+                        for name, module in self.transformer.named_modules():
+                            if hasattr(module, 'scaling'):
+                                if isinstance(module.scaling, dict):
+                                    # Handle ModuleDict case (PEFT implementation)
+                                    if lora_name in module.scaling:
+                                        if isinstance(module.scaling[lora_name], torch.Tensor):
+                                            module.scaling[lora_name] = torch.tensor(
+                                                lora_strength, device=module.scaling[lora_name].device
+                                            )
+                                        else:
+                                            module.scaling[lora_name] = lora_strength
+                                else:
+                                    # Handle direct attribute case for scaling if needed
+                                    if isinstance(module.scaling, torch.Tensor):
+                                        module.scaling = torch.tensor(
+                                            lora_strength, device=module.scaling.device
+                                        )
+                                    else:
+                                        module.scaling = lora_strength
+                else:
+                    print(f"LoRA file for {lora_name} not found!")
+            except Exception as e:
+                print(f"Error loading LoRA {lora_name}: {e}")
+        
+        # Verify LoRA state after loading
+        self.verify_lora_state("After loading LoRAs")

modules/generators/f1_generator.py

@@ -0,0 +1,219 @@
+import torch
+from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
+from diffusers_helper.memory import DynamicSwapInstaller
+from .base_generator import BaseModelGenerator
+
+class F1ModelGenerator(BaseModelGenerator):
+    """
+    Model generator for the F1 HunyuanVideo model.
+    """
+    
+    def __init__(self, **kwargs):
+        """
+        Initialize the F1 model generator.
+        """
+        super().__init__(**kwargs)
+        self.model_name = "F1"
+        self.model_path = 'lllyasviel/FramePack_F1_I2V_HY_20250503'
+    
+    def get_model_name(self):
+        """
+        Get the name of the model.
+        """
+        return self.model_name
+    
+    def load_model(self):
+        """
+        Load the F1 transformer model.
+        """
+        print(f"Loading {self.model_name} Transformer...")
+        
+        # Create the transformer model
+        self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
+            self.model_path, 
+            torch_dtype=torch.bfloat16
+        ).cpu()
+        
+        # Configure the model
+        self.transformer.eval()
+        self.transformer.to(dtype=torch.bfloat16)
+        self.transformer.requires_grad_(False)
+        
+        # Set up dynamic swap if not in high VRAM mode
+        if not self.high_vram:
+            DynamicSwapInstaller.install_model(self.transformer, device=self.gpu)
+        
+        print(f"{self.model_name} Transformer Loaded.")
+        return self.transformer
+    
+    def prepare_history_latents(self, height, width):
+        """
+        Prepare the history latents tensor for the F1 model.
+        
+        Args:
+            height: The height of the image
+            width: The width of the image
+            
+        Returns:
+            The initialized history latents tensor
+        """
+        return torch.zeros(
+            size=(1, 16, 16 + 2 + 1, height // 8, width // 8), 
+            dtype=torch.float32
+        ).cpu()
+    
+    def initialize_with_start_latent(self, history_latents, start_latent):
+        """
+        Initialize the history latents with the start latent for the F1 model.
+        
+        Args:
+            history_latents: The history latents
+            start_latent: The start latent
+            
+        Returns:
+            The initialized history latents
+        """
+        # Add the start frame to history_latents
+        return torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
+    
+    def get_latent_paddings(self, total_latent_sections):
+        """
+        Get the latent paddings for the F1 model.
+        
+        Args:
+            total_latent_sections: The total number of latent sections
+            
+        Returns:
+            A list of latent paddings
+        """
+        # F1 model uses a fixed approach with just 0 for last section and 1 for others
+        return [1] * (total_latent_sections - 1) + [0]
+    
+    def prepare_indices(self, latent_padding_size, latent_window_size):
+        """
+        Prepare the indices for the F1 model.
+        
+        Args:
+            latent_padding_size: The size of the latent padding
+            latent_window_size: The size of the latent window
+            
+        Returns:
+            A tuple of (clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices)
+        """
+        # F1 model uses a different indices approach
+        # latent_window_sizeが4.5の場合は特別に5を使用
+        effective_window_size = 5 if latent_window_size == 4.5 else int(latent_window_size)
+        indices = torch.arange(0, sum([1, 16, 2, 1, latent_window_size])).unsqueeze(0)
+        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)
+        clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
+        
+        return clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices
+    
+    def prepare_clean_latents(self, start_latent, history_latents):
+        """
+        Prepare the clean latents for the F1 model.
+        
+        Args:
+            start_latent: The start latent
+            history_latents: The history latents
+            
+        Returns:
+            A tuple of (clean_latents, clean_latents_2x, clean_latents_4x)
+        """
+        # For F1, we take the last frames for clean latents
+        clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2)
+        # For F1, we prepend the start latent to clean_latents_1x
+        clean_latents = torch.cat([start_latent.to(history_latents), clean_latents_1x], dim=2)
+        
+        return clean_latents, clean_latents_2x, clean_latents_4x
+    
+    def update_history_latents(self, history_latents, generated_latents):
+        """
+        Update the history latents with the generated latents for the F1 model.
+        
+        Args:
+            history_latents: The history latents
+            generated_latents: The generated latents
+            
+        Returns:
+            The updated history latents
+        """
+        # For F1, we append new frames to the end
+        return torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
+    
+    def get_real_history_latents(self, history_latents, total_generated_latent_frames):
+        """
+        Get the real history latents for the F1 model.
+        
+        Args:
+            history_latents: The history latents
+            total_generated_latent_frames: The total number of generated latent frames
+            
+        Returns:
+            The real history latents
+        """
+        # For F1, we take frames from the end
+        return history_latents[:, :, -total_generated_latent_frames:, :, :]
+    
+    def update_history_pixels(self, history_pixels, current_pixels, overlapped_frames):
+        """
+        Update the history pixels with the current pixels for the F1 model.
+        
+        Args:
+            history_pixels: The history pixels
+            current_pixels: The current pixels
+            overlapped_frames: The number of overlapped frames
+            
+        Returns:
+            The updated history pixels
+        """
+        from diffusers_helper.utils import soft_append_bcthw
+        # For F1 model, history_pixels is first, current_pixels is second
+        return soft_append_bcthw(history_pixels, current_pixels, overlapped_frames)
+    
+    def get_section_latent_frames(self, latent_window_size, is_last_section):
+        """
+        Get the number of section latent frames for the F1 model.
+        
+        Args:
+            latent_window_size: The size of the latent window
+            is_last_section: Whether this is the last section
+            
+        Returns:
+            The number of section latent frames
+        """
+        return latent_window_size * 2
+    
+    def get_current_pixels(self, real_history_latents, section_latent_frames, vae):
+        """
+        Get the current pixels for the F1 model.
+        
+        Args:
+            real_history_latents: The real history latents
+            section_latent_frames: The number of section latent frames
+            vae: The VAE model
+            
+        Returns:
+            The current pixels
+        """
+        from diffusers_helper.hunyuan import vae_decode
+        # For F1, we take frames from the end
+        return vae_decode(real_history_latents[:, :, -section_latent_frames:], vae).cpu()
+    
+    def format_position_description(self, total_generated_latent_frames, current_pos, original_pos, current_prompt):
+        """
+        Format the position description for the F1 model.
+        
+        Args:
+            total_generated_latent_frames: The total number of generated latent frames
+            current_pos: The current position in seconds
+            original_pos: The original position in seconds
+            current_prompt: The current prompt
+            
+        Returns:
+            The formatted position description
+        """
+        return (f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, '
+                f'Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 30):.2f} seconds (FPS-30). '
+                f'Current position: {current_pos:.2f}s. '
+                f'using prompt: {current_prompt[:256]}...')

modules/generators/original_generator.py

@@ -0,0 +1,202 @@
+import torch
+from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
+from diffusers_helper.memory import DynamicSwapInstaller
+from .base_generator import BaseModelGenerator
+
+class OriginalModelGenerator(BaseModelGenerator):
+    """
+    Model generator for the Original HunyuanVideo model.
+    """
+    
+    def __init__(self, **kwargs):
+        """
+        Initialize the Original model generator.
+        """
+        super().__init__(**kwargs)
+        self.model_name = "Original"
+        self.model_path = 'lllyasviel/FramePackI2V_HY'
+    
+    def get_model_name(self):
+        """
+        Get the name of the model.
+        """
+        return self.model_name
+    
+    def load_model(self):
+        """
+        Load the Original transformer model.
+        """
+        print(f"Loading {self.model_name} Transformer...")
+        
+        # Create the transformer model
+        self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
+            self.model_path, 
+            torch_dtype=torch.bfloat16
+        ).cpu()
+        
+        # Configure the model
+        self.transformer.eval()
+        self.transformer.to(dtype=torch.bfloat16)
+        self.transformer.requires_grad_(False)
+        
+        # Set up dynamic swap if not in high VRAM mode
+        if not self.high_vram:
+            DynamicSwapInstaller.install_model(self.transformer, device=self.gpu)
+        
+        print(f"{self.model_name} Transformer Loaded.")
+        return self.transformer
+    
+    def prepare_history_latents(self, height, width):
+        """
+        Prepare the history latents tensor for the Original model.
+        
+        Args:
+            height: The height of the image
+            width: The width of the image
+            
+        Returns:
+            The initialized history latents tensor
+        """
+        return torch.zeros(
+            size=(1, 16, 1 + 2 + 16, height // 8, width // 8), 
+            dtype=torch.float32
+        ).cpu()
+    
+    def get_latent_paddings(self, total_latent_sections):
+        """
+        Get the latent paddings for the Original model.
+        
+        Args:
+            total_latent_sections: The total number of latent sections
+            
+        Returns:
+            A list of latent paddings
+        """
+        # Original model uses reversed latent paddings
+        if total_latent_sections > 4:
+            return [3] + [2] * (total_latent_sections - 3) + [1, 0]
+        else:
+            return list(reversed(range(total_latent_sections)))
+    
+    def prepare_indices(self, latent_padding_size, latent_window_size):
+        """
+        Prepare the indices for the Original model.
+        
+        Args:
+            latent_padding_size: The size of the latent padding
+            latent_window_size: The size of the latent window
+            
+        Returns:
+            A tuple of (clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices)
+        """
+        indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0)
+        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)
+        clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1)
+        
+        return clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices
+    
+    def prepare_clean_latents(self, start_latent, history_latents):
+        """
+        Prepare the clean latents for the Original model.
+        
+        Args:
+            start_latent: The start latent
+            history_latents: The history latents
+            
+        Returns:
+            A tuple of (clean_latents, clean_latents_2x, clean_latents_4x)
+        """
+        clean_latents_pre = start_latent.to(history_latents)
+        clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2)
+        clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2)
+        
+        return clean_latents, clean_latents_2x, clean_latents_4x
+    
+    def update_history_latents(self, history_latents, generated_latents):
+        """
+        Update the history latents with the generated latents for the Original model.
+        
+        Args:
+            history_latents: The history latents
+            generated_latents: The generated latents
+            
+        Returns:
+            The updated history latents
+        """
+        # For Original model, we prepend the generated latents
+        return torch.cat([generated_latents.to(history_latents), history_latents], dim=2)
+    
+    def get_real_history_latents(self, history_latents, total_generated_latent_frames):
+        """
+        Get the real history latents for the Original model.
+        
+        Args:
+            history_latents: The history latents
+            total_generated_latent_frames: The total number of generated latent frames
+            
+        Returns:
+            The real history latents
+        """
+        return history_latents[:, :, :total_generated_latent_frames, :, :]
+    
+    def update_history_pixels(self, history_pixels, current_pixels, overlapped_frames):
+        """
+        Update the history pixels with the current pixels for the Original model.
+        
+        Args:
+            history_pixels: The history pixels
+            current_pixels: The current pixels
+            overlapped_frames: The number of overlapped frames
+            
+        Returns:
+            The updated history pixels
+        """
+        from diffusers_helper.utils import soft_append_bcthw
+        # For Original model, current_pixels is first, history_pixels is second
+        return soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
+    
+    def get_section_latent_frames(self, latent_window_size, is_last_section):
+        """
+        Get the number of section latent frames for the Original model.
+        
+        Args:
+            latent_window_size: The size of the latent window
+            is_last_section: Whether this is the last section
+            
+        Returns:
+            The number of section latent frames
+        """
+        return latent_window_size * 2
+    
+    def get_current_pixels(self, real_history_latents, section_latent_frames, vae):
+        """
+        Get the current pixels for the Original model.
+        
+        Args:
+            real_history_latents: The real history latents
+            section_latent_frames: The number of section latent frames
+            vae: The VAE model
+            
+        Returns:
+            The current pixels
+        """
+        from diffusers_helper.hunyuan import vae_decode
+        return vae_decode(real_history_latents[:, :, :section_latent_frames], vae).cpu()
+    
+    def format_position_description(self, total_generated_latent_frames, current_pos, original_pos, current_prompt):
+        """
+        Format the position description for the Original model.
+        
+        Args:
+            total_generated_latent_frames: The total number of generated latent frames
+            current_pos: The current position in seconds
+            original_pos: The original position in seconds
+            current_prompt: The current prompt
+            
+        Returns:
+            The formatted position description
+        """
+        return (f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, '
+                f'Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 30):.2f} seconds (FPS-30). '
+                f'Current position: {current_pos:.2f}s (original: {original_pos:.2f}s). '
+                f'using prompt: {current_prompt[:256]}...')

modules/interface.py

@@ -0,0 +1,771 @@
+import gradio as gr
+import time
+import datetime
+import random
+import json
+import os
+import shutil
+from typing import List, Dict, Any, Optional
+from PIL import Image
+import numpy as np
+import base64
+import io
+
+from modules.video_queue import JobStatus, Job
+from modules.prompt_handler import get_section_boundaries, get_quick_prompts, parse_timestamped_prompt
+from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
+from diffusers_helper.bucket_tools import find_nearest_bucket
+
+def create_interface(
+    process_fn,
+    monitor_fn,
+    end_process_fn,
+    update_queue_status_fn,
+    load_lora_file_fn,
+    job_queue,
+    settings,
+    default_prompt: str = '[1s: The person waves hello] [3s: The person jumps up and down] [5s: The person does a dance]',
+    lora_names: list = [],
+    lora_values: list = []
+):
+    """
+    Create the Gradio interface for the video generation application
+
+    Args:
+        process_fn: Function to process a new job
+        monitor_fn: Function to monitor an existing job
+        end_process_fn: Function to cancel the current job
+        update_queue_status_fn: Function to update the queue status display
+        default_prompt: Default prompt text
+        lora_names: List of loaded LoRA names
+
+    Returns:
+        Gradio Blocks interface
+    """
+    # Get section boundaries and quick prompts
+    section_boundaries = get_section_boundaries()
+    quick_prompts = get_quick_prompts()
+
+    # Create the interface
+    css = make_progress_bar_css()
+    css += """
+    /* Image container styling - more aggressive approach */
+    .contain-image, .contain-image > div, .contain-image > div > img {
+        object-fit: contain !important;
+    }
+    
+    /* Target all images in the contain-image class and its children */
+    .contain-image img,
+    .contain-image > div > img,
+    .contain-image * img {
+        object-fit: contain !important;
+        width: 100% !important;
+        height: 100% !important;
+        max-height: 100% !important;
+        max-width: 100% !important;
+    }
+    
+    /* Additional selectors to override Gradio defaults */
+    .gradio-container img,
+    .gradio-container .svelte-1b5oq5x,
+    .gradio-container [data-testid="image"] img {
+        object-fit: contain !important;
+    }
+    
+    /* Toolbar styling */
+    #fixed-toolbar {
+        position: fixed;
+        top: 0;
+        left: 0;
+        width: 100vw;
+        z-index: 1000;
+        background: rgb(11, 15, 25);
+        color: #fff;
+        padding: 10px 20px;
+        display: flex;
+        align-items: center;
+        gap: 16px;
+        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+        border-bottom: 1px solid #4f46e5;
+    }
+    
+    /* Button styling */
+    #toolbar-add-to-queue-btn button {
+        font-size: 14px !important;
+        padding: 4px 16px !important;
+        height: 32px !important;
+        min-width: 80px !important;
+    }
+    .narrow-button {
+        min-width: 40px !important;
+        width: 40px !important;
+        padding: 0 !important;
+        margin: 0 !important;
+    }
+    .gr-button-primary {
+        color: white;
+    }
+    
+    /* Layout adjustments */
+    body, .gradio-container {
+        padding-top: 40px !important;
+    }
+    """
+
+    # Get the theme from settings
+    current_theme = settings.get("gradio_theme", "default") # Use default if not found
+    block = gr.Blocks(css=css, title="FramePack Studio", theme=current_theme).queue()
+
+    with block:
+
+        with gr.Row(elem_id="fixed-toolbar"):
+            gr.Markdown("<h1 style='margin:0;color:white;'>FramePack Studio</h1>")
+            # with gr.Column(scale=1):
+            #     queue_stats_display = gr.Markdown("<p style='margin:0;color:white;'>Queue: 0 | Completed: 0</p>")
+            with gr.Column(scale=0):
+                refresh_stats_btn = gr.Button("⟳", elem_id="refresh-stats-btn")
+
+
+        
+
+        with gr.Tabs():
+            with gr.Tab("Generate", id="generate_tab"):
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        model_type = gr.Radio(
+                            choices=["Original", "F1"],
+                            value="Original",
+                            label="Model",
+                            info="Select which model to use for generation"
+                        )
+                        input_image = gr.Image(
+                            sources='upload',
+                            type="numpy",
+                            label="Image (optional)",
+                            height=420,
+                            elem_classes="contain-image",
+                            image_mode="RGB",
+                            show_download_button=False,
+                            show_label=True,
+                            container=True
+                        )
+
+                        with gr.Accordion("Latent Image Options", open=False):
+                            latent_type = gr.Dropdown(
+                                ["Black", "White", "Noise", "Green Screen"], label="Latent Image", value="Black", info="Used as a starting point if no image is provided"
+                            )
+
+                        prompt = gr.Textbox(label="Prompt", value=default_prompt)
+
+                        with gr.Accordion("Prompt Parameters", open=False):
+                            n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=True)  # Make visible for both models
+
+                            blend_sections = gr.Slider(
+                                minimum=0, maximum=10, value=4, step=1,
+                                label="Number of sections to blend between prompts"
+                            )
+                        with gr.Accordion("Generation Parameters", open=True):
+                            with gr.Row():
+                                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)
+                                total_second_length = gr.Slider(label="Video Length (Seconds)", minimum=1, maximum=120, value=6, step=0.1)
+                            with gr.Group():
+                                with gr.Row("Resolution"):
+                                    resolutionW = gr.Slider(
+                                        label="Width", minimum=128, maximum=768, value=640, step=32, 
+                                        info="Nearest valid width will be used."
+                                    )
+                                    resolutionH = gr.Slider(
+                                        label="Height", minimum=128, maximum=768, value=640, step=32, 
+                                        info="Nearest valid height will be used."
+                                    )
+                                resolution_text = gr.Markdown(value="<div style='text-align:right; padding:5px 15px 5px 5px;'>Selected bucket for resolution: 640 x 640</div>", label="", show_label=False)
+                            def on_input_image_change(img):
+                                if img is not None:
+                                    return gr.update(info="Nearest valid bucket size will be used. Height will be adjusted automatically."), gr.update(visible=False)
+                                else:
+                                    return gr.update(info="Nearest valid width will be used."), gr.update(visible=True)
+                            input_image.change(fn=on_input_image_change, inputs=[input_image], outputs=[resolutionW, resolutionH])
+                            def on_resolution_change(img, resolutionW, resolutionH):
+                                out_bucket_resH, out_bucket_resW = [640, 640]
+                                if img is not None:
+                                    H, W, _ = img.shape
+                                    out_bucket_resH, out_bucket_resW = find_nearest_bucket(H, W, resolution=resolutionW)
+                                else:
+                                    out_bucket_resH, out_bucket_resW = find_nearest_bucket(resolutionH, resolutionW, (resolutionW+resolutionH)/2) # if resolutionW > resolutionH else resolutionH
+                                return gr.update(value=f"<div style='text-align:right; padding:5px 15px 5px 5px;'>Selected bucket for resolution: {out_bucket_resW} x {out_bucket_resH}</div>")
+                            resolutionW.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
+                            resolutionH.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
+                            with gr.Row("LoRAs"):
+                                lora_selector = gr.Dropdown(
+                                    choices=lora_names,
+                                    label="Select LoRAs to Load",
+                                    multiselect=True,
+                                    value=[],
+                                    info="Select one or more LoRAs to use for this job"
+                                )
+                                lora_names_states = gr.State(lora_names)
+                                lora_sliders = {}
+                                for lora in lora_names:
+                                    lora_sliders[lora] = gr.Slider(
+                                        minimum=0.0, maximum=2.0, value=1.0, step=0.01,
+                                        label=f"{lora} Weight", visible=False, interactive=True
+                                    )
+
+                            with gr.Row("Metadata"):
+                                json_upload = gr.File(
+                                    label="Upload Metadata JSON (optional)",
+                                    file_types=[".json"],
+                                    type="filepath",
+                                    height=100,
+                                )
+                            with gr.Row("TeaCache"):
+                                use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
+
+                            with gr.Row():
+                                seed = gr.Number(label="Seed", value=31337, precision=0)
+                                randomize_seed = gr.Checkbox(label="Randomize", value=False, info="Generate a new random seed for each job")
+
+                        with gr.Accordion("Advanced Parameters", open=False):
+                            latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=True, info='Change at your own risk, very experimental')  # Should not change
+                            cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False)  # Should not change
+                            gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.01)
+                            rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False)  # Should not change
+
+                    with gr.Column():
+                        preview_image = gr.Image(
+                            label="Next Latents", 
+                            height=150, 
+                            visible=True, 
+                            type="numpy", 
+                            interactive=False,
+                            elem_classes="contain-image",
+                            image_mode="RGB"
+                        )
+                        result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=256, loop=True)
+                        progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
+                        progress_bar = gr.HTML('', elem_classes='no-generating-animation')
+
+                        with gr.Row():
+                            current_job_id = gr.Textbox(label="Current Job ID", visible=True, interactive=True)
+                            end_button = gr.Button(value="Cancel Current Job", interactive=True)
+                            start_button = gr.Button(value="Add to Queue", elem_id="toolbar-add-to-queue-btn")
+
+            with gr.Tab("Queue"):
+                with gr.Row():
+                    with gr.Column():
+                        # Create a container for the queue status
+                        with gr.Row():
+                            queue_status = gr.DataFrame(
+                                headers=["Job ID", "Type", "Status", "Created", "Started", "Completed", "Elapsed"], # Removed Preview header
+                                datatype=["str", "str", "str", "str", "str", "str", "str"], # Removed image datatype
+                                label="Job Queue"
+                            )
+                        with gr.Row():
+                            refresh_button = gr.Button("Refresh Queue")
+                            # Connect the refresh button (Moved inside 'with block')
+                            refresh_button.click(
+                                fn=update_queue_status_fn, # Use the function passed in
+                                inputs=[],
+                                outputs=[queue_status]
+                            )
+                        # Create a container for thumbnails (kept for potential future use, though not displayed in DataFrame)
+                        with gr.Row():
+                            thumbnail_container = gr.Column()
+                            thumbnail_container.elem_classes = ["thumbnail-container"]
+
+                        # Add CSS for thumbnails
+            with gr.TabItem("Outputs"):
+                outputDirectory_video = settings.get("output_dir", settings.default_settings['output_dir'])
+                outputDirectory_metadata = settings.get("metadata_dir", settings.default_settings['metadata_dir'])
+                def get_gallery_items():
+                    items = []
+                    for f in os.listdir(outputDirectory_metadata):
+                        if f.endswith(".png"):
+                            prefix = os.path.splitext(f)[0]
+                            latest_video = get_latest_video_version(prefix)
+                            if latest_video:
+                                video_path = os.path.join(outputDirectory_video, latest_video)
+                                mtime = os.path.getmtime(video_path)
+                                preview_path = os.path.join(outputDirectory_metadata, f)
+                                items.append((preview_path, prefix, mtime))
+                    items.sort(key=lambda x: x[2], reverse=True)
+                    return [(i[0], i[1]) for i in items]
+                def get_latest_video_version(prefix):
+                    max_number = -1
+                    selected_file = None
+                    for f in os.listdir(outputDirectory_video):
+                        if f.startswith(prefix + "_") and f.endswith(".mp4"):
+                            num = int(f.replace(prefix + "_", '').replace(".mp4", ''))
+                            if num > max_number:
+                                max_number = num
+                                selected_file = f
+                    return selected_file
+                def load_video_and_info_from_prefix(prefix):
+                    video_file = get_latest_video_version(prefix)
+                    if not video_file:
+                        return None, "JSON not found."
+                    video_path = os.path.join(outputDirectory_video, video_file)
+                    json_path = os.path.join(outputDirectory_metadata, prefix) + ".json"
+                    info = {"description": "no info"}
+                    if os.path.exists(json_path):
+                        with open(json_path, "r", encoding="utf-8") as f:
+                            info = json.load(f)
+                    return video_path, json.dumps(info, indent=2, ensure_ascii=False)
+                gallery_items_state = gr.State(get_gallery_items())
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        thumbs = gr.Gallery(
+                            # value=[i[0] for i in get_gallery_items()],
+                            columns=[4],
+                            allow_preview=False,
+                            object_fit="cover",
+                            height="auto"
+                        )
+                        refresh_button = gr.Button("Update")
+                    with gr.Column(scale=5):
+                        video_out = gr.Video(sources=[], autoplay=True, loop=True, visible=False)
+                    with gr.Column(scale=1):
+                        info_out = gr.Textbox(label="Generation info", visible=False)
+                    def refresh_gallery():
+                        new_items = get_gallery_items()
+                        return gr.update(value=[i[0] for i in new_items]), new_items
+                    refresh_button.click(fn=refresh_gallery, outputs=[thumbs, gallery_items_state])
+                    def on_select(evt: gr.SelectData, gallery_items):
+                        prefix = gallery_items[evt.index][1]
+                        video, info = load_video_and_info_from_prefix(prefix)
+                        return gr.update(value=video, visible=True), gr.update(value=info, visible=True)
+                    thumbs.select(fn=on_select, inputs=[gallery_items_state], outputs=[video_out, info_out])
+            with gr.Tab("Settings"):
+                with gr.Row():
+                    with gr.Column():
+                        save_metadata = gr.Checkbox(
+                            label="Save Metadata", 
+                            info="Save to JSON file", 
+                            value=settings.get("save_metadata", 6),
+                        )
+                        gpu_memory_preservation = gr.Slider(
+                            label="GPU Inference Preserved Memory (GB) (larger means slower)",
+                            minimum=1,
+                            maximum=128,
+                            step=0.1,
+                            value=settings.get("gpu_memory_preservation", 6),
+                            info="Set this number to a larger value if you encounter OOM. Larger value causes slower speed."
+                        )
+                        mp4_crf = gr.Slider(
+                            label="MP4 Compression",
+                            minimum=0,
+                            maximum=100,
+                            step=1,
+                            value=settings.get("mp4_crf", 16),
+                            info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs."
+                        )
+                        clean_up_videos = gr.Checkbox(
+                            label="Clean up video files",
+                            value=settings.get("clean_up_videos", True),
+                            info="If checked, only the final video will be kept after generation."
+                        )
+                        output_dir = gr.Textbox(
+                            label="Output Directory",
+                            value=settings.get("output_dir"),
+                            placeholder="Path to save generated videos"
+                        )
+                        metadata_dir = gr.Textbox(
+                            label="Metadata Directory",
+                            value=settings.get("metadata_dir"),
+                            placeholder="Path to save metadata files"
+                        )
+                        lora_dir = gr.Textbox(
+                            label="LoRA Directory",
+                            value=settings.get("lora_dir"),
+                            placeholder="Path to LoRA models"
+                        )
+                        gradio_temp_dir = gr.Textbox(label="Gradio Temporary Directory", value=settings.get("gradio_temp_dir"))
+                        auto_save = gr.Checkbox(
+                            label="Auto-save settings",
+                            value=settings.get("auto_save_settings", True)
+                        )
+                        # Add Gradio Theme Dropdown
+                        gradio_themes = ["default", "base", "soft", "glass", "mono", "huggingface"]
+                        theme_dropdown = gr.Dropdown(
+                            label="Theme",
+                            choices=gradio_themes,
+                            value=settings.get("gradio_theme", "soft"),
+                            info="Select the Gradio UI theme. Requires restart."
+                        )
+                        save_btn = gr.Button("Save Settings")
+                        cleanup_btn = gr.Button("Clean Up Temporary Files")
+                        status = gr.HTML("")
+                        cleanup_output = gr.Textbox(label="Cleanup Status", interactive=False)
+
+                        def save_settings(save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, selected_theme):
+                            try:
+                                settings.save_settings(
+                                    save_metadata=save_metadata,
+                                    gpu_memory_preservation=gpu_memory_preservation,
+                                    mp4_crf=mp4_crf,
+                                    clean_up_videos=clean_up_videos,
+                                    output_dir=output_dir,
+                                    metadata_dir=metadata_dir,
+                                    lora_dir=lora_dir,
+                                    gradio_temp_dir=gradio_temp_dir,
+                                    auto_save_settings=auto_save,
+                                    gradio_theme=selected_theme
+                                )
+                                return "<p style='color:green;'>Settings saved successfully! Restart required for theme change.</p>"
+                            except Exception as e:
+                                return f"<p style='color:red;'>Error saving settings: {str(e)}</p>"
+
+                        save_btn.click(
+                            fn=save_settings,
+                            inputs=[save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, theme_dropdown],
+                            outputs=[status]
+                        )
+
+                        def cleanup_temp_files():
+                            """Clean up temporary files and folders in the Gradio temp directory"""
+                            temp_dir = settings.get("gradio_temp_dir")
+                            if not temp_dir or not os.path.exists(temp_dir):
+                                return "No temporary directory found or directory does not exist."
+                            
+                            try:
+                                # Get all items in the temp directory
+                                items = os.listdir(temp_dir)
+                                removed_count = 0
+                                print(f"Finding items in {temp_dir}")
+                                for item in items:
+                                    item_path = os.path.join(temp_dir, item)
+                                    try:
+                                        if os.path.isfile(item_path) or os.path.islink(item_path):
+                                            print(f"Removing {item_path}")
+                                            os.remove(item_path)
+                                            removed_count += 1
+                                        elif os.path.isdir(item_path):
+                                            print(f"Removing directory {item_path}")
+                                            shutil.rmtree(item_path)
+                                            removed_count += 1
+                                    except Exception as e:
+                                        print(f"Error removing {item_path}: {e}")
+                                
+                                return f"Cleaned up {removed_count} temporary files/folders."
+                            except Exception as e:
+                                return f"Error cleaning up temporary files: {str(e)}"
+
+        # --- Event Handlers and Connections (Now correctly indented) ---
+
+        # Connect the main process function (wrapper for adding to queue)
+        def process_with_queue_update(model_type, *args):
+            # Extract all arguments (ensure order matches inputs lists)
+            input_image, prompt_text, n_prompt, seed_value, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, randomize_seed_checked, save_metadata_checked, blend_sections, latent_type, clean_up_videos, selected_loras, resolutionW, resolutionH, *lora_args = args
+
+            # DO NOT parse the prompt here. Parsing happens once in the worker.
+
+            # Use the current seed value as is for this job
+            # Call the process function with all arguments
+            # Pass the model_type and the ORIGINAL prompt_text string to the backend process function
+            result = process_fn(model_type, input_image, prompt_text, n_prompt, seed_value, total_second_length, # Pass original prompt_text string
+                            latent_window_size, steps, cfg, gs, rs,
+                            use_teacache, blend_sections, latent_type, clean_up_videos, selected_loras, resolutionW, resolutionH, *lora_args)
+
+            # If randomize_seed is checked, generate a new random seed for the next job
+            new_seed_value = None
+            if randomize_seed_checked:
+                new_seed_value = random.randint(0, 21474)
+                print(f"Generated new seed for next job: {new_seed_value}")
+
+            # If a job ID was created, automatically start monitoring it and update queue
+            if result and result[1]:  # Check if job_id exists in results
+                job_id = result[1]
+                queue_status_data = update_queue_status_fn()
+
+                # Add the new seed value to the results if randomize is checked
+                if new_seed_value is not None:
+                    return [result[0], job_id, result[2], result[3], result[4], result[5], result[6], queue_status_data, new_seed_value]
+                else:
+                    return [result[0], job_id, result[2], result[3], result[4], result[5], result[6], queue_status_data, gr.update()]
+
+            # If no job ID was created, still return the new seed if randomize is checked
+            if new_seed_value is not None:
+                return result + [update_queue_status_fn(), new_seed_value]
+            else:
+                return result + [update_queue_status_fn(), gr.update()]
+
+        # Custom end process function that ensures the queue is updated
+        def end_process_with_update():
+            queue_status_data = end_process_fn()
+            # Make sure to return the queue status data
+            return queue_status_data
+
+        # --- Inputs Lists ---
+        # --- Inputs for Original Model ---
+        ips = [
+            input_image,
+            prompt,
+            n_prompt,
+            seed,
+            total_second_length,
+            latent_window_size,
+            steps,
+            cfg,
+            gs,
+            rs,
+            gpu_memory_preservation,
+            use_teacache,
+            mp4_crf,
+            randomize_seed,
+            save_metadata,
+            blend_sections,
+            latent_type,
+            clean_up_videos,
+            lora_selector,
+            resolutionW,
+            resolutionH,
+            lora_names_states
+        ]
+        # Add LoRA sliders to the input list
+        ips.extend([lora_sliders[lora] for lora in lora_names])
+
+
+        # --- Connect Buttons ---
+        start_button.click(
+            # Pass the selected model type from the radio buttons
+            fn=lambda selected_model, *args: process_with_queue_update(selected_model, *args),
+            inputs=[model_type] + ips,
+            outputs=[result_video, current_job_id, preview_image, progress_desc, progress_bar, start_button, end_button, queue_status, seed]
+        )
+
+        # Connect the end button to cancel the current job and update the queue
+        end_button.click(
+            fn=end_process_with_update,
+            outputs=[queue_status]
+        )
+
+        # --- Connect Monitoring ---
+        # Auto-monitor the current job when job_id changes
+        # Monitor original tab
+        current_job_id.change(
+            fn=monitor_fn,
+            inputs=[current_job_id],
+            outputs=[result_video, current_job_id, preview_image, progress_desc, progress_bar, start_button, end_button]
+        )
+
+        cleanup_btn.click(
+            fn=cleanup_temp_files,
+            outputs=[cleanup_output]
+        )
+
+
+        # --- Connect Queue Refresh ---
+        refresh_stats_btn.click(
+            fn=lambda: update_queue_status_fn(), # Use update_queue_status_fn passed in
+            inputs=None,
+            outputs=[queue_status]  # Removed queue_stats_display from outputs
+        )
+
+        # Set up auto-refresh for queue status (using a timer)
+        refresh_timer = gr.Number(value=0, visible=False)
+        def refresh_timer_fn():
+            """Updates the timer value periodically to trigger queue refresh"""
+            return int(time.time())
+        # This timer seems unused, maybe intended for block.load()? Keeping definition for now.
+        # refresh_timer.change(
+        #     fn=update_queue_status_fn, # Use the function passed in
+        #     outputs=[queue_status] # Update shared queue status display
+        # )
+
+        # --- Connect LoRA UI ---
+        # Function to update slider visibility based on selection
+        def update_lora_sliders(selected_loras):
+            updates = []
+            # Need to handle potential missing keys if lora_names changes dynamically
+            # For now, assume lora_names passed to create_interface is static
+            for lora in lora_names:
+                 updates.append(gr.update(visible=(lora in selected_loras)))
+            # Ensure the output list matches the number of sliders defined
+            num_sliders = len(lora_sliders)
+            return updates[:num_sliders] # Return only updates for existing sliders
+
+        # Connect the dropdown to the sliders
+        lora_selector.change(
+            fn=update_lora_sliders,
+            inputs=[lora_selector],
+            outputs=[lora_sliders[lora] for lora in lora_names] # Assumes lora_sliders keys match lora_names
+        )
+
+
+        # --- Connect Metadata Loading ---
+        # Function to load metadata from JSON file
+        def load_metadata_from_json(json_path):
+            if not json_path:
+                # Return updates for all potentially affected components
+                num_orig_sliders = len(lora_sliders)
+                return [gr.update()] * (2 + num_orig_sliders)
+
+            try:
+                with open(json_path, 'r') as f:
+                    metadata = json.load(f)
+
+                prompt_val = metadata.get('prompt')
+                seed_val = metadata.get('seed')
+
+                # Check for LoRA values in metadata
+                lora_weights = metadata.get('loras', {}) # Changed key to 'loras' based on studio.py worker
+
+                print(f"Loaded metadata from JSON: {json_path}")
+                print(f"Prompt: {prompt_val}, Seed: {seed_val}")
+
+                # Update the UI components
+                updates = [
+                    gr.update(value=prompt_val) if prompt_val else gr.update(),
+                    gr.update(value=seed_val) if seed_val is not None else gr.update()
+                ]
+
+                # Update LoRA sliders if they exist in metadata
+                for lora in lora_names:
+                    if lora in lora_weights:
+                        updates.append(gr.update(value=lora_weights[lora]))
+                    else:
+                        updates.append(gr.update()) # No change if LoRA not in metadata
+
+                # Ensure the number of updates matches the number of outputs
+                num_orig_sliders = len(lora_sliders)
+                return updates[:2 + num_orig_sliders] # Return updates for prompt, seed, and sliders
+
+            except Exception as e:
+                print(f"Error loading metadata: {e}")
+                num_orig_sliders = len(lora_sliders)
+                return [gr.update()] * (2 + num_orig_sliders)
+
+
+        # Connect JSON metadata loader for Original tab
+        json_upload.change(
+            fn=load_metadata_from_json,
+            inputs=[json_upload],
+            outputs=[prompt, seed] + [lora_sliders[lora] for lora in lora_names]
+        )
+
+
+        # --- Helper Functions (defined within create_interface scope if needed by handlers) ---
+        # Function to get queue statistics
+        def get_queue_stats():
+            try:
+                # Get all jobs from the queue
+                jobs = job_queue.get_all_jobs()
+
+                # Count jobs by status
+                status_counts = {
+                    "QUEUED": 0,
+                    "RUNNING": 0,
+                    "COMPLETED": 0,
+                    "FAILED": 0,
+                    "CANCELLED": 0
+                }
+
+                for job in jobs:
+                    if hasattr(job, 'status'):
+                        status = str(job.status) # Use str() for safety
+                        if status in status_counts:
+                            status_counts[status] += 1
+
+                # Format the display text
+                stats_text = f"Queue: {status_counts['QUEUED']} | Running: {status_counts['RUNNING']} | Completed: {status_counts['COMPLETED']} | Failed: {status_counts['FAILED']} | Cancelled: {status_counts['CANCELLED']}"
+
+                return f"<p style='margin:0;color:white;'>{stats_text}</p>"
+
+            except Exception as e:
+                print(f"Error getting queue stats: {e}")
+                return "<p style='margin:0;color:white;'>Error loading queue stats</p>"
+
+        # Add footer with social links
+        with gr.Row(elem_id="footer"):
+            with gr.Column(scale=1):
+                gr.HTML("""
+                <div style="text-align: center; padding: 20px; color: #666;">
+                    <div style="margin-top: 10px;">
+                        <a href="https://patreon.com/Colinu" target="_blank" style="margin: 0 10px; color: #666; text-decoration: none;">
+                            <i class="fab fa-patreon"></i>Support on Patreon
+                        </a>
+                        <a href="https://discord.gg/MtuM7gFJ3V" target="_blank" style="margin: 0 10px; color: #666; text-decoration: none;">
+                            <i class="fab fa-discord"></i> Discord
+                        </a>
+                        <a href="https://github.com/colinurbs/FramePack-Studio" target="_blank" style="margin: 0 10px; color: #666; text-decoration: none;">
+                            <i class="fab fa-github"></i> GitHub
+                        </a>
+                    </div>
+                </div>
+                """)
+
+        # Add CSS for footer
+
+    return block
+
+
+# --- Top-level Helper Functions (Used by Gradio callbacks, must be defined outside create_interface) ---
+
+def format_queue_status(jobs):
+    """Format job data for display in the queue status table"""
+    rows = []
+    for job in jobs:
+        created = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(job.created_at)) if job.created_at else ""
+        started = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(job.started_at)) if job.started_at else ""
+        completed = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(job.completed_at)) if job.completed_at else ""
+
+        # Calculate elapsed time
+        elapsed_time = ""
+        if job.started_at:
+            if job.completed_at:
+                start_datetime = datetime.datetime.fromtimestamp(job.started_at)
+                complete_datetime = datetime.datetime.fromtimestamp(job.completed_at)
+                elapsed_seconds = (complete_datetime - start_datetime).total_seconds()
+                elapsed_time = f"{elapsed_seconds:.2f}s"
+            else:
+                # For running jobs, calculate elapsed time from now
+                start_datetime = datetime.datetime.fromtimestamp(job.started_at)
+                current_datetime = datetime.datetime.now()
+                elapsed_seconds = (current_datetime - start_datetime).total_seconds()
+                elapsed_time = f"{elapsed_seconds:.2f}s (running)"
+
+        # Get generation type from job data
+        generation_type = getattr(job, 'generation_type', 'Original')
+
+        # Removed thumbnail processing
+
+        rows.append([
+            job.id[:6] + '...',
+            generation_type,
+            job.status.value,
+            created,
+            started,
+            completed,
+            elapsed_time
+            # Removed thumbnail from row data
+        ])
+    return rows
+
+# Create the queue status update function (wrapper around format_queue_status)
+def update_queue_status_with_thumbnails(): # Function name is now slightly misleading, but keep for now to avoid breaking clicks
+    # This function is likely called by the refresh button and potentially the timer
+    # It needs access to the job_queue object
+    # Assuming job_queue is accessible globally or passed appropriately
+    # For now, let's assume it's globally accessible as defined in studio.py
+    # If not, this needs adjustment based on how job_queue is managed.
+    try:
+        # Need access to the global job_queue instance from studio.py
+        # This might require restructuring or passing job_queue differently.
+        # For now, assuming it's accessible (this might fail if run standalone)
+        from __main__ import job_queue # Attempt to import from main script scope
+
+        jobs = job_queue.get_all_jobs()
+        for job in jobs:
+            if job.status == JobStatus.PENDING:
+                job.queue_position = job_queue.get_queue_position(job.id)
+
+        if job_queue.current_job:
+            job_queue.current_job.status = JobStatus.RUNNING
+
+        return format_queue_status(jobs)
+    except ImportError:
+        print("Error: Could not import job_queue. Queue status update might fail.")
+        return [] # Return empty list on error
+    except Exception as e:
+        print(f"Error updating queue status: {e}")
+        return []

modules/prompt_handler.py

@@ -0,0 +1,164 @@
+import re
+from dataclasses import dataclass
+from typing import List, Optional
+
+
+@dataclass
+class PromptSection:
+    """Represents a section of the prompt with specific timing information"""
+    prompt: str
+    start_time: float = 0  # in seconds
+    end_time: Optional[float] = None  # in seconds, None means until the end
+
+
+def snap_to_section_boundaries(prompt_sections: List[PromptSection], latent_window_size: int, fps: int = 30) -> List[PromptSection]:
+    """
+    Adjust timestamps to align with model's internal section boundaries
+    
+    Args:
+        prompt_sections: List of PromptSection objects
+        latent_window_size: Size of the latent window used in the model
+        fps: Frames per second (default: 30)
+        
+    Returns:
+        List of PromptSection objects with aligned timestamps
+    """
+    section_duration = (latent_window_size * 4 - 3) / fps  # Duration of one section in seconds
+    
+    aligned_sections = []
+    for section in prompt_sections:
+        # Snap start time to nearest section boundary
+        aligned_start = round(section.start_time / section_duration) * section_duration
+        
+        # Snap end time to nearest section boundary
+        aligned_end = None
+        if section.end_time is not None:
+            aligned_end = round(section.end_time / section_duration) * section_duration
+        
+        # Ensure minimum section length
+        if aligned_end is not None and aligned_end <= aligned_start:
+            aligned_end = aligned_start + section_duration
+            
+        aligned_sections.append(PromptSection(
+            prompt=section.prompt,
+            start_time=aligned_start,
+            end_time=aligned_end
+        ))
+    
+    return aligned_sections
+
+
+def parse_timestamped_prompt(prompt_text: str, total_duration: float, latent_window_size: int = 9, generation_type: str = "Original") -> List[PromptSection]:
+    """
+    Parse a prompt with timestamps in the format [0s-2s: text] or [3s: text]
+    
+    Args:
+        prompt_text: The input prompt text with optional timestamp sections
+        total_duration: Total duration of the video in seconds
+        latent_window_size: Size of the latent window used in the model
+        generation_type: Type of generation ("Original" or "F1")
+        
+    Returns:
+        List of PromptSection objects with timestamps aligned to section boundaries
+        and reversed to account for reverse generation (only for Original type)
+    """
+    # Default prompt for the entire duration if no timestamps are found
+    if "[" not in prompt_text or "]" not in prompt_text:
+        return [PromptSection(prompt=prompt_text.strip())]
+    
+    sections = []
+    # Find all timestamp sections [time: text]
+    timestamp_pattern = r'\[(\d+(?:\.\d+)?s)(?:-(\d+(?:\.\d+)?s))?\s*:\s*(.*?)\]'
+    regular_text = prompt_text
+    
+    for match in re.finditer(timestamp_pattern, prompt_text):
+        start_time_str = match.group(1)
+        end_time_str = match.group(2)
+        section_text = match.group(3).strip()
+        
+        # Convert time strings to seconds
+        start_time = float(start_time_str.rstrip('s'))
+        end_time = float(end_time_str.rstrip('s')) if end_time_str else None
+        
+        sections.append(PromptSection(
+            prompt=section_text,
+            start_time=start_time,
+            end_time=end_time
+        ))
+        
+        # Remove the processed section from regular_text
+        regular_text = regular_text.replace(match.group(0), "")
+    
+    # If there's any text outside of timestamp sections, use it as a default for the entire duration
+    regular_text = regular_text.strip()
+    if regular_text:
+        sections.append(PromptSection(
+            prompt=regular_text,
+            start_time=0,
+            end_time=None
+        ))
+    
+    # Sort sections by start time
+    sections.sort(key=lambda x: x.start_time)
+    
+    # Fill in end times if not specified
+    for i in range(len(sections) - 1):
+        if sections[i].end_time is None:
+            sections[i].end_time = sections[i+1].start_time
+    
+    # Set the last section's end time to the total duration if not specified
+    if sections and sections[-1].end_time is None:
+        sections[-1].end_time = total_duration
+    
+    # Snap timestamps to section boundaries
+    sections = snap_to_section_boundaries(sections, latent_window_size)
+    
+    # Only reverse timestamps for Original generation type
+    if generation_type == "Original":
+        # Now reverse the timestamps to account for reverse generation
+        reversed_sections = []
+        for section in sections:
+            reversed_start = total_duration - section.end_time if section.end_time is not None else 0
+            reversed_end = total_duration - section.start_time
+            reversed_sections.append(PromptSection(
+                prompt=section.prompt,
+                start_time=reversed_start,
+                end_time=reversed_end
+            ))
+        
+        # Sort the reversed sections by start time
+        reversed_sections.sort(key=lambda x: x.start_time)
+        return reversed_sections
+    
+    return sections
+
+
+def get_section_boundaries(latent_window_size: int = 9, count: int = 10) -> str:
+    """
+    Calculate and format section boundaries for UI display
+    
+    Args:
+        latent_window_size: Size of the latent window used in the model
+        count: Number of boundaries to display
+        
+    Returns:
+        Formatted string of section boundaries
+    """
+    section_duration = (latent_window_size * 4 - 3) / 30
+    return ", ".join([f"{i*section_duration:.1f}s" for i in range(count)])
+
+
+def get_quick_prompts() -> List[List[str]]:
+    """
+    Get a list of example timestamped prompts
+    
+    Returns:
+        List of example prompts formatted for Gradio Dataset
+    """
+    prompts = [
+        '[0s: The person waves hello] [2s: The person jumps up and down] [4s: The person does a spin]',
+        '[0s: The person raises both arms slowly] [2s: The person claps hands enthusiastically]',
+        '[0s: Person gives thumbs up] [1.1s: Person smiles and winks] [2.2s: Person shows two thumbs down]',
+        '[0s: Person looks surprised] [1.1s: Person raises arms above head] [2.2s-3.3s: Person puts hands on hips]'
+    ]
+    return [[x] for x in prompts]

modules/settings.py

@@ -0,0 +1,76 @@
+import json
+from pathlib import Path
+from typing import Dict, Any, Optional
+import os
+
+class Settings:
+    def __init__(self):
+        # Get the project root directory (where settings.py is located)
+        project_root = Path(__file__).parent.parent
+        
+        self.settings_file = project_root / ".framepack" / "settings.json"
+        self.settings_file.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Set default paths relative to project root
+        self.default_settings = {
+            "save_metadata": True,
+            "gpu_memory_preservation": 6,
+            "output_dir": str(project_root / "outputs"),
+            "metadata_dir": str(project_root / "outputs"),
+            "lora_dir": str(project_root / "loras"),
+            "gradio_temp_dir": str(project_root / "temp"),
+            "auto_save_settings": True,
+            "gradio_theme": "base",
+            "mp4_crf": 16,
+            "clean_up_videos": True
+        }
+        self.settings = self.load_settings()
+
+    def load_settings(self) -> Dict[str, Any]:
+        """Load settings from file or return defaults"""
+        if self.settings_file.exists():
+            try:
+                with open(self.settings_file, 'r') as f:
+                    loaded_settings = json.load(f)
+                    # Merge with defaults to ensure all settings exist
+                    settings = self.default_settings.copy()
+                    settings.update(loaded_settings)
+                    return settings
+            except Exception as e:
+                print(f"Error loading settings: {e}")
+                return self.default_settings.copy()
+        return self.default_settings.copy()
+
+    def save_settings(self, **kwargs):
+        """Save settings to file. Accepts keyword arguments for any settings to update."""
+        # Update self.settings with any provided keyword arguments
+        self.settings.update(kwargs)
+        # Ensure all default fields are present
+        for k, v in self.default_settings.items():
+            self.settings.setdefault(k, v)
+
+        # Ensure directories exist for relevant fields
+        for dir_key in ["output_dir", "metadata_dir", "lora_dir", "gradio_temp_dir"]:
+            dir_path = self.settings.get(dir_key)
+            if dir_path:
+                os.makedirs(dir_path, exist_ok=True)
+
+        # Save to file
+        with open(self.settings_file, 'w') as f:
+            json.dump(self.settings, f, indent=4)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a setting value"""
+        return self.settings.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a setting value"""
+        self.settings[key] = value
+        if self.settings.get("auto_save_settings", True):
+            self.save_settings()
+
+    def update(self, settings: Dict[str, Any]) -> None:
+        """Update multiple settings at once"""
+        self.settings.update(settings)
+        if self.settings.get("auto_save_settings", True):
+            self.save_settings()

modules/video_queue.py

@@ -0,0 +1,341 @@
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Any, Optional, List
+import queue as queue_module  # Renamed to avoid conflicts
+import io
+import base64
+from PIL import Image
+import numpy as np
+
+from diffusers_helper.thread_utils import AsyncStream
+
+
+# Simple LIFO queue implementation to avoid dependency on queue.LifoQueue
+class SimpleLifoQueue:
+    def __init__(self):
+        self._queue = []
+        self._mutex = threading.Lock()
+        self._not_empty = threading.Condition(self._mutex)
+    
+    def put(self, item):
+        with self._mutex:
+            self._queue.append(item)
+            self._not_empty.notify()
+    
+    def get(self):
+        with self._not_empty:
+            while not self._queue:
+                self._not_empty.wait()
+            return self._queue.pop()
+    
+    def task_done(self):
+        pass  # For compatibility with queue.Queue
+
+
+class JobStatus(Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+@dataclass
+class Job:
+    id: str
+    params: Dict[str, Any]
+    status: JobStatus = JobStatus.PENDING
+    created_at: float = field(default_factory=time.time)
+    started_at: Optional[float] = None
+    completed_at: Optional[float] = None
+    error: Optional[str] = None
+    result: Optional[str] = None
+    progress_data: Optional[Dict] = None
+    queue_position: Optional[int] = None
+    stream: Optional[Any] = None
+    input_image: Optional[np.ndarray] = None
+    latent_type: Optional[str] = None
+    thumbnail: Optional[str] = None
+    generation_type: Optional[str] = None # Added generation_type
+
+    def __post_init__(self):
+        # Store generation type
+        self.generation_type = self.params.get('model_type', 'Original') # Initialize generation_type
+
+        # Store input image or latent type
+        if 'input_image' in self.params and self.params['input_image'] is not None:
+            self.input_image = self.params['input_image']
+            # Create thumbnail
+            img = Image.fromarray(self.input_image)
+            img.thumbnail((100, 100))
+            buffered = io.BytesIO()
+            img.save(buffered, format="PNG")
+            self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+        elif 'latent_type' in self.params:
+            self.latent_type = self.params['latent_type']
+            # Create a colored square based on latent type
+            color_map = {
+                "Black": (0, 0, 0),
+                "White": (255, 255, 255),
+                "Noise": (128, 128, 128),
+                "Green Screen": (0, 177, 64)
+            }
+            color = color_map.get(self.latent_type, (0, 0, 0))
+            img = Image.new('RGB', (100, 100), color)
+            buffered = io.BytesIO()
+            img.save(buffered, format="PNG")
+            self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+
+
+class VideoJobQueue:
+    def __init__(self):
+        self.queue = queue_module.Queue()  # Using standard Queue instead of LifoQueue
+        self.jobs = {}
+        self.current_job = None
+        self.lock = threading.Lock()
+        self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
+        self.worker_thread.start()
+        self.worker_function = None  # Will be set from outside
+        self.is_processing = False  # Flag to track if we're currently processing a job
+    
+    def set_worker_function(self, worker_function):
+        """Set the worker function to use for processing jobs"""
+        self.worker_function = worker_function
+    
+    def add_job(self, params):
+        """Add a job to the queue and return its ID"""
+        job_id = str(uuid.uuid4())
+        job = Job(
+            id=job_id,
+            params=params,
+            status=JobStatus.PENDING,
+            created_at=time.time(),
+            progress_data={},
+            stream=AsyncStream()
+        )
+        
+        with self.lock:
+            print(f"Adding job {job_id} to queue, current job is {self.current_job.id if self.current_job else 'None'}")
+            self.jobs[job_id] = job
+            self.queue.put(job_id)
+        
+        return job_id
+    
+    def get_job(self, job_id):
+        """Get job by ID"""
+        with self.lock:
+            return self.jobs.get(job_id)
+    
+    def get_all_jobs(self):
+        """Get all jobs"""
+        with self.lock:
+            return list(self.jobs.values())
+    
+    def cancel_job(self, job_id):
+        """Cancel a pending job"""
+        with self.lock:
+            job = self.jobs.get(job_id)
+            if job and job.status == JobStatus.PENDING:
+                job.status = JobStatus.CANCELLED
+                job.completed_at = time.time()  # Mark completion time
+                return True
+            elif job and job.status == JobStatus.RUNNING:
+                # Send cancel signal to the job's stream
+                job.stream.input_queue.push('end')
+                # Mark job as cancelled (this will be confirmed when the worker processes the end signal)
+                job.status = JobStatus.CANCELLED
+                job.completed_at = time.time()  # Mark completion time
+                return True
+            return False
+    
+    def get_queue_position(self, job_id):
+        """Get position in queue (0 = currently running)"""
+        with self.lock:
+            job = self.jobs.get(job_id)
+            if not job:
+                return None
+                
+            if job.status == JobStatus.RUNNING:
+                return 0
+                
+            if job.status != JobStatus.PENDING:
+                return None
+                
+            # Count pending jobs ahead in queue
+            position = 1  # Start at 1 because 0 means running
+            for j in self.jobs.values():
+                if (j.status == JobStatus.PENDING and 
+                    j.created_at < job.created_at):
+                    position += 1
+            return position
+    
+    def update_job_progress(self, job_id, progress_data):
+        """Update job progress data"""
+        with self.lock:
+            job = self.jobs.get(job_id)
+            if job:
+                job.progress_data = progress_data
+    
+    def _worker_loop(self):
+        """Worker thread that processes jobs from the queue"""
+        while True:
+            try:
+                # Get the next job ID from the queue
+                try:
+                    job_id = self.queue.get(block=True, timeout=1.0)  # Use timeout to allow periodic checks
+                except queue_module.Empty:
+                    # No jobs in queue, just continue the loop
+                    continue
+                
+                with self.lock:
+                    job = self.jobs.get(job_id)
+                    if not job:
+                        self.queue.task_done()
+                        continue
+                    
+                    # Skip cancelled jobs
+                    if job.status == JobStatus.CANCELLED:
+                        self.queue.task_done()
+                        continue
+                    
+                    # If we're already processing a job, wait for it to complete
+                    if self.is_processing:
+                        # Put the job back in the queue
+                        self.queue.put(job_id)
+                        self.queue.task_done()
+                        time.sleep(0.1)  # Small delay to prevent busy waiting
+                        continue
+                    
+                    print(f"Starting job {job_id}, current job was {self.current_job.id if self.current_job else 'None'}")
+                    job.status = JobStatus.RUNNING
+                    job.started_at = time.time()
+                    self.current_job = job
+                    self.is_processing = True
+                
+                job_completed = False
+                
+                try:
+                    if self.worker_function is None:
+                        raise ValueError("Worker function not set. Call set_worker_function() first.")
+                    
+                    # Start the worker function with the job parameters
+                    from diffusers_helper.thread_utils import async_run
+                    async_run(
+                        self.worker_function,
+                        **job.params,
+                        job_stream=job.stream
+                    )
+                    
+                    # Process the results from the stream
+                    output_filename = None
+                    
+                    # Set a maximum time to wait for the job to complete
+                    max_wait_time = 3600  # 1 hour in seconds
+                    start_time = time.time()
+                    last_activity_time = time.time()
+                    
+                    while True:
+                        # Check if job has been cancelled before processing next output
+                        with self.lock:
+                            if job.status == JobStatus.CANCELLED:
+                                print(f"Job {job_id} was cancelled, breaking out of processing loop")
+                                job_completed = True
+                                break
+                        
+                        # Check if we've been waiting too long without any activity
+                        current_time = time.time()
+                        if current_time - start_time > max_wait_time:
+                            print(f"Job {job_id} timed out after {max_wait_time} seconds")
+                            with self.lock:
+                                job.status = JobStatus.FAILED
+                                job.error = "Job timed out"
+                                job.completed_at = time.time()
+                            job_completed = True
+                            break
+                        
+                        # Check for inactivity (no output for a while)
+                        if current_time - last_activity_time > 60:  # 1 minute of inactivity
+                            print(f"Checking if job {job_id} is still active...")
+                            # Just a periodic check, don't break yet
+                        
+                        try:
+                            # Try to get data from the queue with a non-blocking approach
+                            flag, data = job.stream.output_queue.next()
+                            
+                            # Update activity time since we got some data
+                            last_activity_time = time.time()
+                            
+                            if flag == 'file':
+                                output_filename = data
+                                with self.lock:
+                                    job.result = output_filename
+                            
+                            elif flag == 'progress':
+                                preview, desc, html = data
+                                with self.lock:
+                                    job.progress_data = {
+                                        'preview': preview,
+                                        'desc': desc,
+                                        'html': html
+                                    }
+                            
+                            elif flag == 'end':
+                                print(f"Received end signal for job {job_id}")
+                                job_completed = True
+                                break
+                                
+                        except IndexError:
+                            # Queue is empty, wait a bit and try again
+                            time.sleep(0.1)
+                            continue
+                        except Exception as e:
+                            print(f"Error processing job output: {e}")
+                            # Wait a bit before trying again
+                            time.sleep(0.1)
+                            continue
+                except Exception as e:
+                    import traceback
+                    traceback.print_exc()
+                    print(f"Error processing job {job_id}: {e}")
+                    with self.lock:
+                        job.status = JobStatus.FAILED
+                        job.error = str(e)
+                        job.completed_at = time.time()
+                    job_completed = True
+                
+                finally:
+                    with self.lock:
+                        # Make sure we properly clean up the job state
+                        if job.status == JobStatus.RUNNING:
+                            if job_completed:
+                                job.status = JobStatus.COMPLETED
+                            else:
+                                # Something went wrong but we didn't mark it as completed
+                                job.status = JobStatus.FAILED
+                                job.error = "Job processing was interrupted"
+                            
+                            job.completed_at = time.time()
+                        
+                        print(f"Finishing job {job_id} with status {job.status}")
+                        self.is_processing = False
+                        self.current_job = None
+                        self.queue.task_done()
+                
+            except Exception as e:
+                import traceback
+                traceback.print_exc()
+                print(f"Error in worker loop: {e}")
+                
+                # Make sure we reset processing state if there was an error
+                with self.lock:
+                    self.is_processing = False
+                    if self.current_job:
+                        self.current_job.status = JobStatus.FAILED
+                        self.current_job.error = f"Worker loop error: {str(e)}"
+                        self.current_job.completed_at = time.time()
+                        self.current_job = None
+                
+                time.sleep(0.5)  # Prevent tight loop on error

requirements.txt

@@ -0,0 +1,17 @@
+accelerate==1.6.0
+diffusers==0.33.1
+transformers==4.46.2
+gradio==5.25.2
+sentencepiece==0.2.0
+pillow==11.1.0
+av==12.1.0
+numpy==1.26.2
+scipy==1.12.0
+requests==2.31.0
+torchsde==0.2.6
+jinja2>=3.1.2
+torchvision
+einops
+opencv-contrib-python
+safetensors
+peft

studio.py

@@ -0,0 +1,1012 @@
+from diffusers_helper.hf_login import login
+
+import json
+import os
+from pathlib import PurePath
+import time
+import argparse
+import traceback
+import einops
+import numpy as np
+import torch
+import datetime
+
+os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
+
+import gradio as gr
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+from diffusers import AutoencoderKLHunyuanVideo
+from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
+from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
+from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, generate_timestamp
+from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
+from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
+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
+from diffusers_helper.thread_utils import AsyncStream
+from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+from transformers import SiglipImageProcessor, SiglipVisionModel
+from diffusers_helper.clip_vision import hf_clip_vision_encode
+from diffusers_helper.bucket_tools import find_nearest_bucket
+from diffusers_helper import lora_utils
+from diffusers_helper.lora_utils import load_lora, unload_all_loras
+
+# Import model generators
+from modules.generators import create_model_generator
+
+# Global cache for prompt embeddings
+prompt_embedding_cache = {}
+# Import from modules
+from modules.video_queue import VideoJobQueue, JobStatus
+from modules.prompt_handler import parse_timestamped_prompt
+from modules.interface import create_interface, format_queue_status
+from modules.settings import Settings
+
+# ADDED: Debug function to verify LoRA state
+def verify_lora_state(transformer, label=""):
+    """Debug function to verify the state of LoRAs in a transformer model"""
+    if transformer is None:
+        print(f"[{label}] Transformer is None, cannot verify LoRA state")
+        return
+        
+    has_loras = False
+    if hasattr(transformer, 'peft_config'):
+        adapter_names = list(transformer.peft_config.keys()) if transformer.peft_config else []
+        if adapter_names:
+            has_loras = True
+            print(f"[{label}] Transformer has LoRAs: {', '.join(adapter_names)}")
+        else:
+            print(f"[{label}] Transformer has no LoRAs in peft_config")
+    else:
+        print(f"[{label}] Transformer has no peft_config attribute")
+        
+    # Check for any LoRA modules
+    for name, module in transformer.named_modules():
+        if hasattr(module, 'lora_A') and module.lora_A:
+            has_loras = True
+            # print(f"[{label}] Found lora_A in module {name}")
+        if hasattr(module, 'lora_B') and module.lora_B:
+            has_loras = True
+            # print(f"[{label}] Found lora_B in module {name}")
+            
+    if not has_loras:
+        print(f"[{label}] No LoRA components found in transformer")
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--share', action='store_true')
+parser.add_argument("--server", type=str, default='0.0.0.0')
+parser.add_argument("--port", type=int, required=False)
+parser.add_argument("--inbrowser", action='store_true')
+parser.add_argument("--lora", type=str, default=None, help="Lora path (comma separated for multiple)")
+args = parser.parse_args()
+
+print(args)
+
+free_mem_gb = get_cuda_free_memory_gb(gpu)
+high_vram = free_mem_gb > 60
+
+print(f'Free VRAM {free_mem_gb} GB')
+print(f'High-VRAM Mode: {high_vram}')
+
+# Load models
+text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
+text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
+tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
+tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
+vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
+
+feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
+image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
+
+# Initialize model generator placeholder
+current_generator = None # Will hold the currently active model generator
+
+# Load models based on VRAM availability later
+ 
+# Configure models
+vae.eval()
+text_encoder.eval()
+text_encoder_2.eval()
+image_encoder.eval()
+
+if not high_vram:
+   vae.enable_slicing()
+   vae.enable_tiling()
+
+
+vae.to(dtype=torch.float16)
+image_encoder.to(dtype=torch.float16)
+text_encoder.to(dtype=torch.float16)
+text_encoder_2.to(dtype=torch.float16)
+
+vae.requires_grad_(False)
+text_encoder.requires_grad_(False)
+text_encoder_2.requires_grad_(False)
+image_encoder.requires_grad_(False)
+
+# Create lora directory if it doesn't exist
+lora_dir = os.path.join(os.path.dirname(__file__), 'loras')
+os.makedirs(lora_dir, exist_ok=True)
+
+# Initialize LoRA support - moved scanning after settings load
+lora_names = []
+lora_values = [] # This seems unused for population, might be related to weights later
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+
+# Define default LoRA folder path relative to the script directory (used if setting is missing)
+default_lora_folder = os.path.join(script_dir, "loras")
+os.makedirs(default_lora_folder, exist_ok=True) # Ensure default exists
+
+if not high_vram:
+    # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
+    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
+else:
+    text_encoder.to(gpu)
+    text_encoder_2.to(gpu)
+    image_encoder.to(gpu)
+    vae.to(gpu)
+
+stream = AsyncStream()
+
+outputs_folder = './outputs/'
+os.makedirs(outputs_folder, exist_ok=True)
+
+# Initialize settings
+settings = Settings()
+
+# --- Populate LoRA names AFTER settings are loaded ---
+lora_folder_from_settings: str = settings.get("lora_dir", default_lora_folder) # Use setting, fallback to default
+print(f"Scanning for LoRAs in: {lora_folder_from_settings}")
+if os.path.isdir(lora_folder_from_settings):
+    try:
+        lora_files = [f for f in os.listdir(lora_folder_from_settings)
+                     if f.endswith('.safetensors') or f.endswith('.pt')]
+        for lora_file in lora_files:
+            lora_name = PurePath(lora_file).stem
+            lora_names.append(lora_name) # Get name without extension
+        print(f"Found LoRAs: {lora_names}")
+    except Exception as e:
+        print(f"Error scanning LoRA directory '{lora_folder_from_settings}': {e}")
+else:
+    print(f"LoRA directory not found: {lora_folder_from_settings}")
+# --- End LoRA population ---
+
+
+# Create job queue
+job_queue = VideoJobQueue()
+
+
+
+# Function to load a LoRA file
+def load_lora_file(lora_file: str | PurePath):
+    if not lora_file:
+        return None, "No file selected"
+    
+    try:
+        # Get the filename from the path
+        lora_path = PurePath(lora_file)
+        lora_name = lora_path.name
+        
+        # Copy the file to the lora directory
+        lora_dest = PurePath(lora_dir, lora_path)
+        import shutil
+        shutil.copy(lora_file, lora_dest)
+        
+        # Load the LoRA
+        global current_generator, lora_names
+        if current_generator is None:
+            return None, "Error: No model loaded to apply LoRA to. Generate something first."
+        
+        # Unload any existing LoRAs first
+        current_generator.unload_loras()
+        
+        # Load the single LoRA
+        selected_loras = [lora_path.stem]
+        current_generator.load_loras(selected_loras, lora_dir, selected_loras)
+        
+        # Add to lora_names if not already there
+        lora_base_name = lora_path.stem
+        if lora_base_name not in lora_names:
+            lora_names.append(lora_base_name)
+        
+        # Get the current device of the transformer
+        device = next(current_generator.transformer.parameters()).device
+        
+        # Move all LoRA adapters to the same device as the base model
+        current_generator.move_lora_adapters_to_device(device)
+        
+        print(f"Loaded LoRA: {lora_name} to {current_generator.get_model_name()} model")
+        
+        return gr.update(choices=lora_names), f"Successfully loaded LoRA: {lora_name}"
+    except Exception as e:
+        print(f"Error loading LoRA: {e}")
+        return None, f"Error loading LoRA: {e}"
+
+@torch.no_grad()
+def get_cached_or_encode_prompt(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, target_device):
+    """
+    Retrieves prompt embeddings from cache or encodes them if not found.
+    Stores encoded embeddings (on CPU) in the cache.
+    Returns embeddings moved to the target_device.
+    """
+    if prompt in prompt_embedding_cache:
+        print(f"Cache hit for prompt: {prompt[:60]}...")
+        llama_vec_cpu, llama_mask_cpu, clip_l_pooler_cpu = prompt_embedding_cache[prompt]
+        # Move cached embeddings (from CPU) to the target device
+        llama_vec = llama_vec_cpu.to(target_device)
+        llama_attention_mask = llama_mask_cpu.to(target_device) if llama_mask_cpu is not None else None
+        clip_l_pooler = clip_l_pooler_cpu.to(target_device)
+        return llama_vec, llama_attention_mask, clip_l_pooler
+    else:
+        print(f"Cache miss for prompt: {prompt[:60]}...")
+        llama_vec, clip_l_pooler = encode_prompt_conds(
+            prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2
+        )
+        llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
+        # Store CPU copies in cache
+        prompt_embedding_cache[prompt] = (llama_vec.cpu(), llama_attention_mask.cpu() if llama_attention_mask is not None else None, clip_l_pooler.cpu())
+        # Return embeddings already on the target device (as encode_prompt_conds uses the model's device)
+        return llama_vec, llama_attention_mask, clip_l_pooler
+        
+@torch.no_grad()
+def worker(
+    model_type,
+    input_image,
+    prompt_text, 
+    n_prompt, 
+    seed, 
+    total_second_length, 
+    latent_window_size,
+    steps, 
+    cfg, 
+    gs, 
+    rs, 
+    use_teacache, 
+    blend_sections, 
+    latent_type,
+    selected_loras,
+    has_input_image,
+    lora_values=None, 
+    job_stream=None,
+    output_dir=None,
+    metadata_dir=None,
+    resolutionW=640,  # Add resolution parameter with default value
+    resolutionH=640,
+    lora_loaded_names=[]
+):
+    global high_vram, current_generator
+    
+    # Ensure any existing LoRAs are unloaded from the current generator
+    if current_generator is not None:
+        print("Unloading any existing LoRAs before starting new job")
+        current_generator.unload_loras()
+        import gc
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    
+    stream_to_use = job_stream if job_stream is not None else stream
+
+    total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
+    total_latent_sections = int(max(round(total_latent_sections), 1))
+
+    # --- Total progress tracking ---
+    total_steps = total_latent_sections * steps  # Total diffusion steps over all segments
+    step_durations = []  # Rolling history of recent step durations for ETA
+    last_step_time = time.time()
+
+    # Parse the timestamped prompt with boundary snapping and reversing
+    # prompt_text should now be the original string from the job queue
+    prompt_sections = parse_timestamped_prompt(prompt_text, total_second_length, latent_window_size, model_type)
+    job_id = generate_timestamp()
+
+    stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
+
+    try:
+        if not high_vram:
+            # Unload everything *except* the potentially active transformer
+            unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae)
+            if current_generator is not None and current_generator.transformer is not None:
+                offload_model_from_device_for_memory_preservation(current_generator.transformer, target_device=gpu, preserved_memory_gb=8)
+
+        # --- Model Loading / Switching ---
+        print(f"Worker starting for model type: {model_type}")
+        
+        # Create the appropriate model generator
+        new_generator = create_model_generator(
+            model_type,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            vae=vae,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            high_vram=high_vram,
+            prompt_embedding_cache=prompt_embedding_cache,
+            settings=settings
+        )
+        
+        # Update the global generator
+        current_generator = new_generator
+        
+        # Load the transformer model
+        current_generator.load_model()
+        
+        # Ensure the model has no LoRAs loaded
+        print(f"Ensuring {model_type} model has no LoRAs loaded")
+        current_generator.unload_loras()
+
+        # Pre-encode all prompts
+        stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding all prompts...'))))
+
+        if not high_vram:
+            fake_diffusers_current_device(text_encoder, gpu)
+            load_model_as_complete(text_encoder_2, target_device=gpu)
+
+        # PROMPT BLENDING: Pre-encode all prompts and store in a list in order
+        unique_prompts = []
+        for section in prompt_sections:
+            if section.prompt not in unique_prompts:
+                unique_prompts.append(section.prompt)
+
+        encoded_prompts = {}
+        for prompt in unique_prompts:
+            # Use the helper function for caching and encoding
+            llama_vec, llama_attention_mask, clip_l_pooler = get_cached_or_encode_prompt(
+                prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, gpu
+            )
+            encoded_prompts[prompt] = (llama_vec, llama_attention_mask, clip_l_pooler)
+
+        # PROMPT BLENDING: Build a list of (start_section_idx, prompt) for each prompt
+        prompt_change_indices = []
+        last_prompt = None
+        for idx, section in enumerate(prompt_sections):
+            if section.prompt != last_prompt:
+                prompt_change_indices.append((idx, section.prompt))
+                last_prompt = section.prompt
+
+        # Encode negative prompt
+        if cfg == 1:
+            llama_vec_n, llama_attention_mask_n, clip_l_pooler_n = (
+                torch.zeros_like(encoded_prompts[prompt_sections[0].prompt][0]),
+                torch.zeros_like(encoded_prompts[prompt_sections[0].prompt][1]),
+                torch.zeros_like(encoded_prompts[prompt_sections[0].prompt][2])
+            )
+        else:
+             # Use the helper function for caching and encoding negative prompt
+            llama_vec_n, llama_attention_mask_n, clip_l_pooler_n = get_cached_or_encode_prompt(
+                n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, gpu
+            )
+
+        # Processing input image
+        stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
+
+        H, W, _ = input_image.shape
+        height, width = find_nearest_bucket(H, W, resolution=resolutionW if has_input_image else (resolutionH+resolutionW)/2)
+        input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
+
+        if settings.get("save_metadata"):
+            metadata = PngInfo()
+            # prompt_text should be a string here now
+            metadata.add_text("prompt", prompt_text)
+            metadata.add_text("seed", str(seed))
+            Image.fromarray(input_image_np).save(os.path.join(metadata_dir, f'{job_id}.png'), pnginfo=metadata)
+
+            metadata_dict = {
+                "prompt": prompt_text, # Use the original string
+                "seed": seed,
+                "total_second_length": total_second_length,
+                "steps": steps,
+                "cfg": cfg,
+                "gs": gs,
+                "rs": rs,
+                "latent_type" : latent_type,
+                "blend_sections": blend_sections,
+                "latent_window_size": latent_window_size,
+                "mp4_crf": settings.get("mp4_crf"),
+                "timestamp": time.time(),
+                "resolutionW": resolutionW,  # Add resolution to metadata
+                "resolutionH": resolutionH,
+                "model_type": model_type  # Add model type to metadata
+            }
+            # Add LoRA information to metadata if LoRAs are used
+            def ensure_list(x):
+                if isinstance(x, list):
+                    return x
+                elif x is None:
+                    return []
+                else:
+                    return [x]
+
+            selected_loras = ensure_list(selected_loras)
+            lora_values = ensure_list(lora_values)
+
+            if selected_loras and len(selected_loras) > 0:
+                lora_data = {}
+                for lora_name in selected_loras:
+                    try:
+                        idx = lora_loaded_names.index(lora_name)
+                        weight = lora_values[idx] if lora_values and idx < len(lora_values) else 1.0
+                        if isinstance(weight, list):
+                            weight_value = weight[0] if weight and len(weight) > 0 else 1.0
+                        else:
+                            weight_value = weight
+                        lora_data[lora_name] = float(weight_value)
+                    except ValueError:
+                        lora_data[lora_name] = 1.0
+                metadata_dict["loras"] = lora_data
+
+            with open(os.path.join(metadata_dir, f'{job_id}.json'), 'w') as f:
+                json.dump(metadata_dict, f, indent=2)
+        else:
+            Image.fromarray(input_image_np).save(os.path.join(metadata_dir, f'{job_id}.png'))
+
+        input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
+        input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
+
+        # VAE encoding
+        stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
+
+        if not high_vram:
+            load_model_as_complete(vae, target_device=gpu)
+
+        start_latent = vae_encode(input_image_pt, vae)
+
+        # CLIP Vision
+        stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
+
+        if not high_vram:
+            load_model_as_complete(image_encoder, target_device=gpu)
+
+        image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
+        image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
+
+        # Dtype
+        for prompt_key in encoded_prompts:
+            llama_vec, llama_attention_mask, clip_l_pooler = encoded_prompts[prompt_key]
+            llama_vec = llama_vec.to(current_generator.transformer.dtype)
+            clip_l_pooler = clip_l_pooler.to(current_generator.transformer.dtype)
+            encoded_prompts[prompt_key] = (llama_vec, llama_attention_mask, clip_l_pooler)
+
+        llama_vec_n = llama_vec_n.to(current_generator.transformer.dtype)
+        clip_l_pooler_n = clip_l_pooler_n.to(current_generator.transformer.dtype)
+        image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(current_generator.transformer.dtype)
+
+        # Sampling
+        stream_to_use.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
+
+        rnd = torch.Generator("cpu").manual_seed(seed)
+        num_frames = latent_window_size * 4 - 3
+
+        # Initialize history latents based on model type
+        history_latents = current_generator.prepare_history_latents(height, width)
+        
+        # For F1 model, initialize with start latent
+        if model_type == "F1":
+            history_latents = current_generator.initialize_with_start_latent(history_latents, start_latent)
+            total_generated_latent_frames = 1  # Start with 1 for F1 model since it includes the first frame
+
+        history_pixels = None
+        if model_type == "Original":
+            total_generated_latent_frames = 0
+        
+        # Get latent paddings from the generator
+        latent_paddings = current_generator.get_latent_paddings(total_latent_sections)
+
+        # PROMPT BLENDING: Track section index
+        section_idx = 0
+
+        # Load LoRAs if selected
+        if selected_loras:
+            current_generator.load_loras(selected_loras, lora_folder_from_settings, lora_loaded_names, lora_values)
+
+        # --- Callback for progress ---
+        def callback(d):
+            nonlocal last_step_time, step_durations
+            now_time = time.time()
+            # Record duration between diffusion steps (skip first where duration may include setup)
+            if last_step_time is not None:
+                step_delta = now_time - last_step_time
+                if step_delta > 0:
+                    step_durations.append(step_delta)
+                    if len(step_durations) > 30:  # Keep only recent 30 steps
+                        step_durations.pop(0)
+            last_step_time = now_time
+            avg_step = sum(step_durations) / len(step_durations) if step_durations else 0.0
+
+            preview = d['denoised']
+            preview = vae_decode_fake(preview)
+            preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
+            preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
+
+            # --- Progress & ETA logic ---
+            # Current segment progress
+            current_step = d['i'] + 1
+            percentage = int(100.0 * current_step / steps)
+
+            # Total progress
+            total_steps_done = section_idx * steps + current_step
+            total_percentage = int(100.0 * total_steps_done / total_steps)
+
+            # ETA calculations
+            def fmt_eta(sec):
+                try:
+                    return str(datetime.timedelta(seconds=int(sec)))
+                except Exception:
+                    return "--:--"
+
+            segment_eta = (steps - current_step) * avg_step if avg_step else 0
+            total_eta = (total_steps - total_steps_done) * avg_step if avg_step else 0
+
+            segment_hint = f'Sampling {current_step}/{steps}  ETA {fmt_eta(segment_eta)}'
+            total_hint = f'Total {total_steps_done}/{total_steps}  ETA {fmt_eta(total_eta)}'
+
+            current_pos = (total_generated_latent_frames * 4 - 3) / 30
+            original_pos = total_second_length - current_pos
+            if current_pos < 0: current_pos = 0
+            if original_pos < 0: original_pos = 0
+
+            hint = segment_hint  # deprecated variable kept to minimise other code changes
+            desc = current_generator.format_position_description(
+                total_generated_latent_frames, 
+                current_pos, 
+                original_pos, 
+                current_prompt
+            )
+
+            progress_data = {
+                'preview': preview,
+                'desc': desc,
+                'html': make_progress_bar_html(percentage, segment_hint) + make_progress_bar_html(total_percentage, total_hint)
+            }
+            if job_stream is not None:
+                job = job_queue.get_job(job_id)
+                if job:
+                    job.progress_data = progress_data
+
+            stream_to_use.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, segment_hint) + make_progress_bar_html(total_percentage, total_hint))))
+
+        # --- Main generation loop ---
+        for latent_padding in latent_paddings:
+            is_last_section = latent_padding == 0
+            latent_padding_size = latent_padding * latent_window_size
+
+            if stream_to_use.input_queue.top() == 'end':
+                stream_to_use.output_queue.push(('end', None))
+                return
+
+            current_time_position = (total_generated_latent_frames * 4 - 3) / 30  # in seconds
+            if current_time_position < 0:
+                current_time_position = 0.01
+
+            # Find the appropriate prompt for this section
+            current_prompt = prompt_sections[0].prompt  # Default to first prompt
+            for section in prompt_sections:
+                if section.start_time <= current_time_position and (section.end_time is None or current_time_position < section.end_time):
+                    current_prompt = section.prompt
+                    break
+
+            # PROMPT BLENDING: Find if we're in a blend window
+            blend_alpha = None
+            prev_prompt = current_prompt
+            next_prompt = current_prompt
+
+            # Only try to blend if we have prompt change indices and multiple sections
+            if prompt_change_indices and len(prompt_sections) > 1:
+                for i, (change_idx, prompt) in enumerate(prompt_change_indices):
+                    if section_idx < change_idx:
+                        prev_prompt = prompt_change_indices[i - 1][1] if i > 0 else prompt
+                        next_prompt = prompt
+                        blend_start = change_idx
+                        blend_end = change_idx + blend_sections
+                        if section_idx >= change_idx and section_idx < blend_end:
+                            blend_alpha = (section_idx - change_idx + 1) / blend_sections
+                        break
+                    elif section_idx == change_idx:
+                        # At the exact change, start blending
+                        if i > 0:
+                            prev_prompt = prompt_change_indices[i - 1][1]
+                            next_prompt = prompt
+                            blend_alpha = 1.0 / blend_sections
+                        else:
+                            prev_prompt = prompt
+                            next_prompt = prompt
+                            blend_alpha = None
+                        break
+                else:
+                    # After last change, no blending
+                    prev_prompt = current_prompt
+                    next_prompt = current_prompt
+                    blend_alpha = None
+
+            # Get the encoded prompt for this section
+            if blend_alpha is not None and prev_prompt != next_prompt:
+                # Blend embeddings
+                prev_llama_vec, prev_llama_attention_mask, prev_clip_l_pooler = encoded_prompts[prev_prompt]
+                next_llama_vec, next_llama_attention_mask, next_clip_l_pooler = encoded_prompts[next_prompt]
+                llama_vec = (1 - blend_alpha) * prev_llama_vec + blend_alpha * next_llama_vec
+                llama_attention_mask = prev_llama_attention_mask  # usually same
+                clip_l_pooler = (1 - blend_alpha) * prev_clip_l_pooler + blend_alpha * next_clip_l_pooler
+                print(f"Blending prompts: '{prev_prompt[:30]}...' -> '{next_prompt[:30]}...', alpha={blend_alpha:.2f}")
+            else:
+                llama_vec, llama_attention_mask, clip_l_pooler = encoded_prompts[current_prompt]
+
+            original_time_position = total_second_length - current_time_position
+            if original_time_position < 0:
+                original_time_position = 0
+
+            print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, '
+                  f'time position: {current_time_position:.2f}s (original: {original_time_position:.2f}s), '
+                  f'using prompt: {current_prompt[:60]}...')
+
+            # Prepare indices using the generator
+            clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices = current_generator.prepare_indices(latent_padding_size, latent_window_size)
+
+            # Prepare clean latents using the generator
+            clean_latents, clean_latents_2x, clean_latents_4x = current_generator.prepare_clean_latents(start_latent, history_latents)
+            
+            # Print debug info
+            print(f"{model_type} model section {section_idx+1}/{total_latent_sections}, latent_padding={latent_padding}")
+
+            if not high_vram:
+                # Unload VAE etc. before loading transformer
+                unload_complete_models(vae, text_encoder, text_encoder_2, image_encoder)
+                move_model_to_device_with_memory_preservation(current_generator.transformer, target_device=gpu, preserved_memory_gb=settings.get("gpu_memory_preservation"))
+                if selected_loras:
+                    current_generator.move_lora_adapters_to_device(gpu)
+
+            if use_teacache:
+                current_generator.transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
+            else:
+                current_generator.transformer.initialize_teacache(enable_teacache=False)
+
+            generated_latents = sample_hunyuan(
+                transformer=current_generator.transformer,
+                sampler='unipc',
+                width=width,
+                height=height,
+                frames=num_frames,
+                real_guidance_scale=cfg,
+                distilled_guidance_scale=gs,
+                guidance_rescale=rs,
+                num_inference_steps=steps,
+                generator=rnd,
+                prompt_embeds=llama_vec,
+                prompt_embeds_mask=llama_attention_mask,
+                prompt_poolers=clip_l_pooler,
+                negative_prompt_embeds=llama_vec_n,
+                negative_prompt_embeds_mask=llama_attention_mask_n,
+                negative_prompt_poolers=clip_l_pooler_n,
+                device=gpu,
+                dtype=torch.bfloat16,
+                image_embeddings=image_encoder_last_hidden_state,
+                latent_indices=latent_indices,
+                clean_latents=clean_latents,
+                clean_latent_indices=clean_latent_indices,
+                clean_latents_2x=clean_latents_2x,
+                clean_latent_2x_indices=clean_latent_2x_indices,
+                clean_latents_4x=clean_latents_4x,
+                clean_latent_4x_indices=clean_latent_4x_indices,
+                callback=callback,
+            )
+
+            total_generated_latent_frames += int(generated_latents.shape[2])
+            # Update history latents using the generator
+            history_latents = current_generator.update_history_latents(history_latents, generated_latents)
+
+            if not high_vram:
+                if selected_loras:
+                    current_generator.move_lora_adapters_to_device(cpu)
+                offload_model_from_device_for_memory_preservation(current_generator.transformer, target_device=gpu, preserved_memory_gb=8)
+                load_model_as_complete(vae, target_device=gpu)
+
+            # Get real history latents using the generator
+            real_history_latents = current_generator.get_real_history_latents(history_latents, total_generated_latent_frames)
+
+            if history_pixels is None:
+                history_pixels = vae_decode(real_history_latents, vae).cpu()
+            else:
+                section_latent_frames = current_generator.get_section_latent_frames(latent_window_size, is_last_section)
+                overlapped_frames = latent_window_size * 4 - 3
+
+                # Get current pixels using the generator
+                current_pixels = current_generator.get_current_pixels(real_history_latents, section_latent_frames, vae)
+                
+                # Update history pixels using the generator
+                history_pixels = current_generator.update_history_pixels(history_pixels, current_pixels, overlapped_frames)
+                
+                print(f"{model_type} model section {section_idx+1}/{total_latent_sections}, history_pixels shape: {history_pixels.shape}")
+
+            if not high_vram:
+                unload_complete_models()
+
+            output_filename = os.path.join(output_dir, f'{job_id}_{total_generated_latent_frames}.mp4')
+            save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=settings.get("mp4_crf"))
+            print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
+            stream_to_use.output_queue.push(('file', output_filename))
+
+            if is_last_section:
+                break
+
+            section_idx += 1  # PROMPT BLENDING: increment section index
+
+        # Unload all LoRAs after generation completed
+        if selected_loras:
+            print("Unloading all LoRAs after generation completed")
+            current_generator.unload_loras()
+            import gc
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    except:
+        traceback.print_exc()
+        # Unload all LoRAs after error
+        if current_generator is not None and selected_loras:
+            print("Unloading all LoRAs after error")
+            current_generator.unload_loras()
+            import gc
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
+        stream_to_use.output_queue.push(('error', f"Error during generation: {traceback.format_exc()}"))
+        if not high_vram:
+            # Ensure all models including the potentially active transformer are unloaded on error
+            unload_complete_models(
+                text_encoder, text_encoder_2, image_encoder, vae, 
+                current_generator.transformer if current_generator else None
+            )
+
+    if settings.get("clean_up_videos"):
+        try:
+            video_files = [
+                f for f in os.listdir(output_dir)
+                if f.startswith(f"{job_id}_") and f.endswith(".mp4")
+            ]
+            print(f"Video files found for cleanup: {video_files}")
+            if video_files:
+                def get_frame_count(filename):
+                    try:
+                        # Handles filenames like jobid_123.mp4
+                        return int(filename.replace(f"{job_id}_", "").replace(".mp4", ""))
+                    except Exception:
+                        return -1
+                video_files_sorted = sorted(video_files, key=get_frame_count)
+                print(f"Sorted video files: {video_files_sorted}")
+                final_video = video_files_sorted[-1]
+                for vf in video_files_sorted[:-1]:
+                    full_path = os.path.join(output_dir, vf)
+                    try:
+                        os.remove(full_path)
+                        print(f"Deleted intermediate video: {full_path}")
+                    except Exception as e:
+                        print(f"Failed to delete {full_path}: {e}")
+        except Exception as e:
+            print(f"Error during video cleanup: {e}")
+
+    # Final verification of LoRA state
+    if current_generator and current_generator.transformer:
+        verify_lora_state(current_generator.transformer, "Worker end")
+
+    stream_to_use.output_queue.push(('end', None))
+    return
+
+
+
+# Set the worker function for the job queue
+job_queue.set_worker_function(worker)
+
+
+def process(
+        model_type,
+        input_image,
+        prompt_text,
+        n_prompt,
+        seed, 
+        total_second_length, 
+        latent_window_size, 
+        steps, 
+        cfg, 
+        gs, 
+        rs, 
+        use_teacache, 
+        blend_sections, 
+        latent_type,
+        clean_up_videos,
+        selected_loras,
+        resolutionW,
+        resolutionH,
+        lora_loaded_names,
+        *lora_values
+    ):
+    
+    # Create a blank black image if no 
+    # Create a default image based on the selected latent_type
+    has_input_image = True
+    if input_image is None:
+        has_input_image = False
+        default_height, default_width = resolutionH, resolutionW
+        if latent_type == "White":
+            # Create a white image
+            input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
+            print("No input image provided. Using a blank white image.")
+
+        elif latent_type == "Noise":
+            # Create a noise image
+            input_image = np.random.randint(0, 256, (default_height, default_width, 3), dtype=np.uint8)
+            print("No input image provided. Using a random noise image.")
+
+        elif latent_type == "Green Screen":
+            # Create a green screen image with standard chroma key green (0, 177, 64)
+            input_image = np.zeros((default_height, default_width, 3), dtype=np.uint8)
+            input_image[:, :, 1] = 177  # Green channel
+            input_image[:, :, 2] = 64   # Blue channel
+            # Red channel remains 0
+            print("No input image provided. Using a standard chroma key green screen.")
+
+        else:  # Default to "Black" or any other value
+            # Create a black image
+            input_image = np.zeros((default_height, default_width, 3), dtype=np.uint8)
+            print(f"No input image provided. Using a blank black image (latent_type: {latent_type}).")
+
+    
+    # Create job parameters
+    job_params = {
+        'model_type': model_type,
+        'input_image': input_image.copy(),  # Make a copy to avoid reference issues
+        'prompt_text': prompt_text,
+        'n_prompt': n_prompt,
+        'seed': seed,
+        'total_second_length': total_second_length,
+        'latent_window_size': latent_window_size,
+        'latent_type': latent_type,
+        'steps': steps,
+        'cfg': cfg,
+        'gs': gs,
+        'rs': rs,
+        'blend_sections': blend_sections,
+        'use_teacache': use_teacache,
+        'selected_loras': selected_loras,
+        'has_input_image': has_input_image,
+        'output_dir': settings.get("output_dir"),
+        'metadata_dir': settings.get("metadata_dir"),
+        'resolutionW': resolutionW, # Add resolution parameter
+        'resolutionH': resolutionH,
+        'lora_loaded_names': lora_loaded_names
+    }
+    
+    # Add LoRA values if provided - extract them from the tuple
+    if lora_values:
+        # Convert tuple to list
+        lora_values_list = list(lora_values)
+        job_params['lora_values'] = lora_values_list
+    
+    # Add job to queue
+    job_id = job_queue.add_job(job_params)
+    
+    # Set the generation_type attribute on the job object directly
+    job = job_queue.get_job(job_id)
+    if job:
+        job.generation_type = model_type  # Set generation_type to model_type for display in queue
+    print(f"Added job {job_id} to queue")
+    
+    queue_status = update_queue_status()
+    # Return immediately after adding to queue
+    return None, job_id, None, '', f'Job added to queue. Job ID: {job_id}', gr.update(interactive=True), gr.update(interactive=True)
+
+
+
+def end_process():
+    """Cancel the current running job and update the queue status"""
+    print("Cancelling current job")
+    with job_queue.lock:
+        if job_queue.current_job:
+            job_id = job_queue.current_job.id
+            print(f"Cancelling job {job_id}")
+
+            # Send the end signal to the job's stream
+            if job_queue.current_job.stream:
+                job_queue.current_job.stream.input_queue.push('end')
+                
+            # Mark the job as cancelled
+            job_queue.current_job.status = JobStatus.CANCELLED
+            job_queue.current_job.completed_at = time.time()  # Set completion time
+    
+    # Force an update to the queue status
+    return update_queue_status()
+
+
+def update_queue_status():
+    """Update queue status and refresh job positions"""
+    jobs = job_queue.get_all_jobs()
+    for job in jobs:
+        if job.status == JobStatus.PENDING:
+            job.queue_position = job_queue.get_queue_position(job.id)
+    
+    # Make sure to update current running job info
+    if job_queue.current_job:
+        # Make sure the running job is showing status = RUNNING
+        job_queue.current_job.status = JobStatus.RUNNING
+    
+    return format_queue_status(jobs)
+
+
+def monitor_job(job_id):
+    """
+    Monitor a specific job and update the UI with the latest video segment as soon as it's available.
+    """
+    if not job_id:
+        yield None, None, None, '', 'No job ID provided', gr.update(interactive=True), gr.update(interactive=True)
+        return
+
+    last_video = None  # Track the last video file shown
+
+    while True:
+        job = job_queue.get_job(job_id)
+        if not job:
+            yield None, job_id, None, '', 'Job not found', gr.update(interactive=True), gr.update(interactive=True)
+            return
+
+        # If a new video file is available, yield it immediately
+        if job.result and job.result != last_video:
+            last_video = job.result
+            # You can also update preview/progress here if desired
+            yield last_video, job_id, gr.update(visible=True), '', '', gr.update(interactive=True), gr.update(interactive=True)
+
+        # Handle job status and progress
+        if job.status == JobStatus.PENDING:
+            position = job_queue.get_queue_position(job_id)
+            yield last_video, job_id, gr.update(visible=True), '', f'Waiting in queue. Position: {position}', gr.update(interactive=True), gr.update(interactive=True)
+
+        elif job.status == JobStatus.RUNNING:
+            if job.progress_data and 'preview' in job.progress_data:
+                preview = job.progress_data.get('preview')
+                desc = job.progress_data.get('desc', '')
+                html = job.progress_data.get('html', '')
+                yield last_video, job_id, gr.update(visible=True, value=preview), desc, html, gr.update(interactive=True), gr.update(interactive=True)
+            else:
+                yield last_video, job_id, gr.update(visible=True), '', 'Processing...', gr.update(interactive=True), gr.update(interactive=True)
+
+        elif job.status == JobStatus.COMPLETED:
+            # Show the final video
+            yield last_video, job_id, gr.update(visible=True), '', '', gr.update(interactive=True), gr.update(interactive=True)
+            break
+
+        elif job.status == JobStatus.FAILED:
+            yield last_video, job_id, gr.update(visible=True), '', f'Error: {job.error}', gr.update(interactive=True), gr.update(interactive=True)
+            break
+
+        elif job.status == JobStatus.CANCELLED:
+            yield last_video, job_id, gr.update(visible=True), '', 'Job cancelled', gr.update(interactive=True), gr.update(interactive=True)
+            break
+
+        # Wait a bit before checking again
+        time.sleep(0.5)
+
+
+# Set Gradio temporary directory from settings
+os.environ["GRADIO_TEMP_DIR"] = settings.get("gradio_temp_dir")
+
+# Create the interface
+interface = create_interface(
+    process_fn=process,
+    monitor_fn=monitor_job,
+    end_process_fn=end_process,
+    update_queue_status_fn=update_queue_status,
+    load_lora_file_fn=load_lora_file,
+    job_queue=job_queue,
+    settings=settings,
+    lora_names=lora_names # Explicitly pass the found LoRA names
+)
+
+# Launch the interface
+interface.launch(
+    server_name=args.server,
+    server_port=args.port,
+    share=args.share,
+    inbrowser=args.inbrowser
+)