Migrated from GitHub

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

ORIGINAL_README.md

@@ -0,0 +1,67 @@
+# goan: A Power-User UI for FramePack
+
+Welcome to `goan`, an enhanced user interface designed for creative professionals and power users of FramePack. This project builds upon the brilliant FramePack video generation engine created by lllyasviel (of Fooocus and the Stable Diffusion Forge fork of A1111), exposing useful controls to reach further into FramePack's functionality through a robust and intuitive interface.
+
+The base FramePack provides a powerful core model. `goan` extends it with a suite of tools designed for serious workflow, experimentation, and reproducibility. Unlock fine-grained control over your video generations with batch processing, parameter editing, effortless recipe sharing, and complete workspace management.
+
+---
+
+### Key Features for a Creative Workflow
+
+The features in `goan` are designed to directly support two important aspects for any serious use of FramePack: deep control over the diffusion process and recovery handling for long-running jobs.
+
+* **Uninterrupted Sessions & True Resilience:** Long-running jobs are the norm for video generation, and nothing is more frustrating than a UI disconnect or a browser crash wiping out hours of progress. `goan` is architected specifically to combat this fragility, a common pain point in many Gradio-based UIs.
+    * **Backend-Driven State:** The entire state of your session—the task queue, processing status, and live updates—is managed persistently on the backend. This means your job keeps running safely on the server, completely independent of the browser tab's status. Whether the tab is minimized, the screen is locked, or the connection is temporarily lost, your progress is secure.
+    * **Intelligent Session Re-attachment:** If your browser tab becomes disconnected for any reason, `goan`'s UI is designed to intelligently and automatically reconnect to the active backend process upon being re-focused. It finds the existing session and its live update stream, allowing the UI to seamlessly catch up to the real-time status of your render queue. This completely mitigates the dreaded "Error" boxes that plague typical Gradio interfaces during long sessions.
+    * **Robust Task Queue & Crash Recovery:** Beyond network stability, the entire task queue is automatically saved to disk. If there's a system crash or you need to restart, you can simply relaunch `goan`, and your queue will be right where you left it, ready to continue processing. No more lost work.
+
+* **Limit MP4 Preview Generation:** The base FramePack functionality writes a VAE-decoded `.mp4` preview for every single segment. This creates a tremendous amount of potentially unneeded compute, causing video generation to take much longer than necessary. `goan` introduces the ability to restrict these previews to only the segments you care about, using a combination of a periodic slider and a comma-separated list of individual segments to dramatically speed up your workflow.
+
+* **Advanced Diffusion Controls:** The base FramePack is a unique approach to video generation. `goan` exposes advanced controls like **Variable CFG**, which allows you to change the prompt adherence over the course of the video. This can be used to correct for a tendency for FramePack to "burn in" or oversaturate the final video as total length increases, giving you greater artistic control.
+
+* **Effortless "Recipe" Sharing:** `goan`'s **Drop-in Parameter Loading** allows you to save all creative settings directly into a generated PNG. Share the image, and anyone using `goan` can drop it into their UI to instantly load your exact "recipe," making collaboration and experimentation simple and repeatable.
+
+* **Complete Workspace Management:** For more complex projects, you can save your *entire UI state*—every slider, checkbox, and text prompt—into a single `.json` file. This ensures you can always get back to a specific setup for consistent results.
+
+### Deeper Dive: Functional & UI Control Comparison
+
+For those curious about the specifics, here's a more detailed breakdown of what's new.
+
+#### Diffusion Controls (CFG, Guidance Scale)
+
+**Understanding CFG:** Classifier-Free Guidance (CFG) is a critical technique in diffusion models. Think of it as a knob that controls how strongly the model should adhere to your text prompt versus how much creative freedom it has.
+* A **low CFG** value allows the model to be more imaginative, potentially straying from the prompt.
+* A **high CFG** value forces the model to follow the prompt more strictly, which can sometimes reduce creativity or lead to artifacts if pushed too high.
+
+In this model, there are two main guidance controls:
+
+* **`Distilled CFG Scale` (`gs`):** This is the primary control you will use.
+    * Recommended settings to begin with are:
+    * Always start at 10.
+    * If bright colors contrast too harshly, experiment with levels around 7-9.
+    * Variable CFG has been added but useful suggestions here are still pending.
+* **`CFG Scale` (`cfg`):** This controls the standard guidance, which is essential for negative prompts.
+    * **Important:** For your **Negative Prompt** to have any effect, you must set the `CFG Scale` to a value greater than 1.0 (e.g., 1.1 or 2.0).
+    * **Performance Trade-Off:** Be aware that setting `CFG Scale` to any value other than 1.0 will roughly **double the generation time** for your video, as it requires a second pass for each step. Use it only when you need the control of a negative prompt.
+
+**Comparison:**
+* **Base FramePack:** Presented a very simplified interface. Key controls like `CFG Scale` and `CFG Re-Scale` were hidden (`visible=False`), and the `info` text for `Distilled CFG Scale` explicitly said, "Changing this value is not recommended." This was effective for a simple demo but limited experimentation.
+* **`goan`:** Exposes all guidance controls for the power user. It introduces the concept of **Variable CFG**, allowing the `Distilled CFG Scale` to change linearly over the course of the generation. This provides advanced control over the video's evolution, letting a user start with high prompt adherence and gradually decrease it, for example.
+
+#### New Functionality: Workspace & Metadata
+
+This entire feature set is new in `goan` to extend power user functionality to FramePack.
+
+* **Drop-in Parameter Loading:** This is the core of the new workflow. You can take a PNG generated by `goan`, drop it into the image input, and the UI will automatically detect the embedded settings. A modal will ask if you want to apply them. This makes sharing and reusing "recipes" effortless.
+* **Workspace Management:** Users can now save the *entire state* of the UI—all sliders, text boxes, and checkboxes—to a `.json` file. This "workspace" can be reloaded at any time, which is invaluable for complex projects or for ensuring consistent settings across sessions.
+* **Session Persistence & Autosave:** The task queue is automatically saved when the application is closed and reloaded on startup. This prevents the loss of a long list of batched jobs. The UI also attempts to restore its last state after a page refresh.
+* **Full Task Queue Control:** `goan` includes a full-featured task queue where you can add, remove, reorder, and *edit* jobs before you start processing. This is a massive improvement over the original's single-task processing model.
+
+### For Developers: A Look Under the Hood (pending review)
+
+### Code Example: A Glimpse at the New Architecture (pending review)
+
+### Acknowledgements
+
+* **@lllyasviel** for creating the groundbreaking FramePack engine.
+* **@Tophness** for the super-useful queueing system architecture introduced in FramePack PR #150, from which forms the foundation of goan's task management, background processing, and progress update features.

diffusers_helper/bucket_tools.py

@@ -0,0 +1,30 @@
+bucket_options = {
+    640: [
+        (416, 960),
+        (448, 864),
+        (480, 832),
+        (512, 768),
+        (544, 704),
+        (576, 672),
+        (608, 640),
+        (640, 608),
+        (672, 576),
+        (704, 544),
+        (768, 512),
+        (832, 480),
+        (864, 448),
+        (960, 416),
+    ],
+}
+
+
+def find_nearest_bucket(h, w, resolution=640):
+    min_metric = float('inf')
+    best_bucket = None
+    for (bucket_h, bucket_w) in bucket_options[resolution]:
+        metric = abs(h * bucket_w - w * bucket_h)
+        if metric <= min_metric:
+            min_metric = metric
+            best_bucket = (bucket_h, bucket_w)
+    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/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,1035 @@
+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
+
+    B, L, H, C = q.shape
+
+    q = q.flatten(0, 1)
+    k = k.flatten(0, 1)
+    v = v.flatten(0, 1)
+
+    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.unflatten(0, (B, L))
+
+    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)
+
+        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)

generation_core.py

@@ -0,0 +1,615 @@
+import torch
+import traceback
+import einops
+import numpy as np
+import os
+import threading
+import json
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+
+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.pipelines.k_diffusion_hunyuan import sample_hunyuan
+from diffusers_helper.memory import (
+    unload_complete_models,
+    load_model_as_complete,
+    move_model_to_device_with_memory_preservation,
+    offload_model_from_device_for_memory_preservation,
+    fake_diffusers_current_device,
+    gpu,
+)
+from diffusers_helper.clip_vision import hf_clip_vision_encode
+from diffusers_helper.bucket_tools import find_nearest_bucket
+from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+from ui import metadata as metadata_manager
+
+
+@torch.no_grad()
+def worker(
+    # --- Task I/O & Identity ---
+    task_id,
+    input_image,
+    output_folder,
+    output_queue_ref,
+    # --- Creative Parameters (The "Recipe") ---
+    prompt,
+    n_prompt,
+    seed,
+    total_second_length,
+    steps,
+    cfg,
+    gs,
+    gs_final,
+    gs_schedule_active,
+    rs,
+    preview_frequency,
+    segments_to_decode_csv,
+    # --- Environment & Debug Parameters ---
+    latent_window_size,
+    gpu_memory_preservation,
+    use_teacache,
+    use_fp32_transformer_output,
+    mp4_crf,
+    # --- Model & System Objects (Passed from main app) ---
+    text_encoder,
+    text_encoder_2,
+    tokenizer,
+    tokenizer_2,
+    vae,
+    feature_extractor,
+    image_encoder,
+    transformer,
+    high_vram,
+    # --- Control Flow ---
+    abort_event: threading.Event = None,
+):
+    outputs_folder = (
+        os.path.expanduser(output_folder) if output_folder else "./outputs/"
+    )
+    os.makedirs(outputs_folder, exist_ok=True)
+
+    # --- Gemini: do not touch - "secret sauce"
+    total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
+    total_latent_sections = int(max(round(total_latent_sections), 1))
+
+    job_id = f"{generate_timestamp()}_task{task_id}"
+    output_queue_ref.push(
+        (
+            "progress",
+            (
+                task_id,
+                None,
+                f"Total Segments: {total_latent_sections}",
+                make_progress_bar_html(0, "Starting ..."),
+            ),
+        )
+    )
+    # ---
+    parsed_segments_to_decode_set = set()
+    if segments_to_decode_csv:
+        try:
+            parsed_segments_to_decode_set = {
+                int(s.strip()) for s in segments_to_decode_csv.split(",") if s.strip()
+            }
+        except ValueError:
+            print(
+                f"Task {task_id}: Warning - Could not parse 'Segments to Decode CSV': \"{segments_to_decode_csv}\"."
+            )
+    final_output_filename = None
+    success = False
+    initial_gs_from_ui = gs
+    gs_final_value_for_schedule = (
+        gs_final if gs_final is not None else initial_gs_from_ui
+    )
+    original_fp32_setting = transformer.high_quality_fp32_output_for_inference
+    transformer.high_quality_fp32_output_for_inference = use_fp32_transformer_output
+    print(
+        f"Task {task_id}: transformer.high_quality_fp32_output_for_inference set to {use_fp32_transformer_output}"
+    )
+
+    try:
+        if not isinstance(input_image, np.ndarray):
+            raise ValueError(f"Task {task_id}: input_image is not a NumPy array.")
+
+        output_queue_ref.push(
+            (
+                "progress",
+                (
+                    task_id,
+                    None,
+                    f"Total Segments: {total_latent_sections}",
+                    make_progress_bar_html(0, "Image processing ..."),
+                ),
+            )
+        )
+        if input_image.shape[-1] == 4:
+            pil_img = Image.fromarray(input_image)
+            input_image = np.array(pil_img.convert("RGB"))
+        H, W, C = input_image.shape
+        if C != 3:
+            raise ValueError(
+                f"Task {task_id}: Input image must be RGB, found {C} channels."
+            )
+        height, width = find_nearest_bucket(H, W, resolution=640)
+        input_image_np = resize_and_center_crop(
+            input_image, target_width=width, target_height=height
+        )
+
+        metadata_obj = PngInfo()
+        params_to_save_in_metadata = {
+            "prompt": prompt,
+            "n_prompt": n_prompt,
+            "seed": seed,
+            "total_second_length": total_second_length,
+            "steps": steps,
+            "cfg": cfg,
+            "gs": gs,
+            "gs_final": gs_final,
+            "gs_schedule_active": gs_schedule_active,
+            "rs": rs,
+            "preview_frequency": preview_frequency,
+            "segments_to_decode_csv": segments_to_decode_csv,
+        }
+        metadata_obj.add_text("parameters", json.dumps(params_to_save_in_metadata))
+        initial_image_with_params_path = os.path.join(
+            outputs_folder, f"{job_id}_initial_image_with_params.png"
+        )
+        try:
+            Image.fromarray(input_image_np).save(
+                initial_image_with_params_path, pnginfo=metadata_obj
+            )
+        except Exception as e_png:
+            print(
+                f"Task {task_id}: WARNING - Failed to save initial image with parameters: {e_png}"
+            )
+
+        # --- Gemini: do not touch - "secret sauce"
+        if not high_vram:
+            unload_complete_models(
+                text_encoder, text_encoder_2, image_encoder, vae, transformer
+            )
+        output_queue_ref.push(
+            (
+                "progress",
+                (
+                    task_id,
+                    None,
+                    f"Total Segments: {total_latent_sections}",
+                    make_progress_bar_html(0, "Text encoding ..."),
+                ),
+            )
+        )
+        if not high_vram:
+            fake_diffusers_current_device(text_encoder, gpu)
+            load_model_as_complete(text_encoder_2, target_device=gpu)
+        llama_vec, clip_l_pooler = encode_prompt_conds(
+            prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2
+        )
+        if cfg == 1:
+            llama_vec_n, clip_l_pooler_n = torch.zeros_like(
+                llama_vec
+            ), torch.zeros_like(clip_l_pooler)
+        else:
+            llama_vec_n, clip_l_pooler_n = encode_prompt_conds(
+                n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2
+            )
+        llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
+        llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(
+            llama_vec_n, length=512
+        )
+        input_image_pt = (
+            torch.from_numpy(input_image_np).float().permute(2, 0, 1).unsqueeze(0)
+            / 127.5
+            - 1.0
+        )
+        input_image_pt = input_image_pt[:, :, None, :, :]
+        output_queue_ref.push(
+            (
+                "progress",
+                (
+                    task_id,
+                    None,
+                    f"Total Segments: {total_latent_sections}",
+                    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)
+        output_queue_ref.push(
+            (
+                "progress",
+                (
+                    task_id,
+                    None,
+                    f"Total Segments: {total_latent_sections}",
+                    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
+        (
+            llama_vec,
+            llama_vec_n,
+            clip_l_pooler,
+            clip_l_pooler_n,
+            image_encoder_last_hidden_state,
+        ) = [
+            t.to(transformer.dtype)
+            for t in [
+                llama_vec,
+                llama_vec_n,
+                clip_l_pooler,
+                clip_l_pooler_n,
+                image_encoder_last_hidden_state,
+            ]
+        ]
+
+        output_queue_ref.push(
+            (
+                "progress",
+                (
+                    task_id,
+                    None,
+                    f"Total Segments: {total_latent_sections}",
+                    make_progress_bar_html(0, "Start sampling ..."),
+                ),
+            )
+        )
+        rnd = torch.Generator(device="cpu").manual_seed(int(seed))
+        num_frames = latent_window_size * 4 - 3
+        # overlapped_frames = num_frames
+
+        history_latents = torch.zeros(
+            size=(1, 16, 1 + 2 + 16, height // 8, width // 8),
+            dtype=torch.float32,
+            device="cpu",
+        )
+        history_pixels = None
+        total_generated_latent_frames = 0
+        latent_paddings = list(reversed(range(total_latent_sections)))
+        if total_latent_sections > 4:
+            latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
+
+        # for latent_padding_iteration, latent_padding in enumerate(latent_paddings):
+        #     if abort_event and abort_event.is_set(): raise KeyboardInterrupt("Abort signal received.")
+        #     is_last_section = (latent_padding == 0)
+        #     latent_padding_size = latent_padding * latent_window_size
+        #     print(f'Task {task_id}: Seg {latent_padding_iteration + 1}/{total_latent_sections} (lp_val={latent_padding}), last_loop_seg={is_last_section}')
+
+        # ^ our code | v Flash code
+
+        for latent_padding_iteration, latent_padding in enumerate(latent_paddings):
+            if abort_event and abort_event.is_set():
+                raise KeyboardInterrupt("Abort signal received.")
+            is_last_section = latent_padding == 0
+            latent_padding_size = latent_padding * latent_window_size
+            # Added for consistent 1-indexed segment number for loop segments
+            current_loop_segment_number = latent_padding_iteration + 1
+            print(
+                f"Task {task_id}: Seg {current_loop_segment_number}/{total_latent_sections} (lp_val={latent_padding}), last_loop_seg={is_last_section}"
+            )
+
+            indices = torch.arange(
+                0,
+                sum([1, latent_padding_size, latent_window_size, 1, 2, 16]),
+                device="cpu",
+            ).unsqueeze(0)
+            (
+                clean_latent_indices_pre,
+                _,
+                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_latents_pre = start_latent.to(
+                history_latents.device, dtype=history_latents.dtype
+            )
+            clean_latent_indices = torch.cat(
+                [clean_latent_indices_pre, clean_latent_indices_post], dim=1
+            )
+            # current_history_depth_for_clean_split = history_latents.shape[2]; needed_depth_for_clean_split = 1 + 2 + 16
+            # history_latents_for_clean_split = history_latents
+            # if current_history_depth_for_clean_split < needed_depth_for_clean_split:
+            #     padding_needed = needed_depth_for_clean_split - current_history_depth_for_clean_split
+            #     pad_tensor = torch.zeros(history_latents.shape[0], history_latents.shape[1], padding_needed, history_latents.shape[3], history_latents.shape[4], dtype=history_latents.dtype, device=history_latents.device)
+            #     history_latents_for_clean_split = torch.cat((history_latents, pad_tensor), dim=2)
+            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)
+
+            if not high_vram:
+                unload_complete_models()
+                move_model_to_device_with_memory_preservation(
+                    transformer,
+                    target_device=gpu,
+                    preserved_memory_gb=gpu_memory_preservation,
+                )
+            transformer.initialize_teacache(
+                enable_teacache=use_teacache, num_steps=steps
+            )
+
+            def callback_diffusion_step(d):
+                if abort_event and abort_event.is_set():
+                    raise KeyboardInterrupt("Abort signal received during sampling.")
+                current_diffusion_step = d["i"] + 1
+                is_first_step = current_diffusion_step == 1
+                is_last_step = current_diffusion_step == steps
+                is_preview_step = preview_frequency > 0 and (
+                    current_diffusion_step % preview_frequency == 0
+                )
+                if not (is_first_step or is_last_step or is_preview_step):
+                    return
+                preview_latent = d["denoised"]
+                preview_img_np = vae_decode_fake(preview_latent)
+                preview_img_np = (
+                    (preview_img_np * 255.0)
+                    .detach()
+                    .cpu()
+                    .numpy()
+                    .clip(0, 255)
+                    .astype(np.uint8)
+                )
+                preview_img_np = einops.rearrange(
+                    preview_img_np, "b c t h w -> (b h) (t w) c"
+                )
+
+                # percentage = int(100.0 * current_diffusion_step / steps)
+                # hint = f'Segment {latent_padding_iteration + 1}, Sampling {current_diffusion_step}/{steps}'
+                # current_video_frames_count = history_pixels.shape[2] if history_pixels is not None else 0
+                # desc = f'Task {task_id}: Vid Frames: {current_video_frames_count}, Len: {current_video_frames_count / 30 :.2f}s. Seg {latent_padding_iteration + 1}/{total_latent_sections}. Extending...'
+                # output_queue_ref.push(('progress', (task_id, preview_img_np, desc, make_progress_bar_html(percentage, hint))))
+
+                # ^ our code | v Flash code
+
+                percentage = int(100.0 * current_diffusion_step / steps)
+                hint = f"Segment {current_loop_segment_number}, Sampling {current_diffusion_step}/{steps}"  # Updated hint
+                current_video_frames_count = (
+                    history_pixels.shape[2] if history_pixels is not None else 0
+                )
+                desc = f"Task {task_id}: Vid Frames: {current_video_frames_count}, Len: {current_video_frames_count / 30 :.2f}s. Seg {current_loop_segment_number}/{total_latent_sections}. Extending..."  # Updated desc
+                output_queue_ref.push(
+                    (
+                        "progress",
+                        (
+                            task_id,
+                            preview_img_np,
+                            desc,
+                            make_progress_bar_html(percentage, hint),
+                        ),
+                    )
+                )
+
+            current_segment_gs_to_use = initial_gs_from_ui
+            if gs_schedule_active and total_latent_sections > 1:
+                progress_for_gs = (
+                    latent_padding_iteration / (total_latent_sections - 1)
+                    if total_latent_sections > 1
+                    else 0
+                )
+                current_segment_gs_to_use = (
+                    initial_gs_from_ui
+                    + (gs_final_value_for_schedule - initial_gs_from_ui)
+                    * progress_for_gs
+                )
+
+            generated_latents = sample_hunyuan(
+                transformer=transformer,
+                sampler="unipc",
+                width=width,
+                height=height,
+                frames=num_frames,
+                real_guidance_scale=cfg,
+                distilled_guidance_scale=current_segment_gs_to_use,
+                guidance_rescale=rs,
+                num_inference_steps=steps,
+                generator=rnd,
+                prompt_embeds=llama_vec.to(transformer.device),
+                prompt_embeds_mask=llama_attention_mask.to(transformer.device),
+                prompt_poolers=clip_l_pooler.to(transformer.device),
+                negative_prompt_embeds=llama_vec_n.to(transformer.device),
+                negative_prompt_embeds_mask=llama_attention_mask_n.to(
+                    transformer.device
+                ),
+                negative_prompt_poolers=clip_l_pooler_n.to(transformer.device),
+                device=transformer.device,
+                dtype=transformer.dtype,
+                image_embeddings=image_encoder_last_hidden_state.to(transformer.device),
+                latent_indices=latent_indices.to(transformer.device),
+                clean_latents=clean_latents.to(
+                    transformer.device, dtype=transformer.dtype
+                ),
+                clean_latent_indices=clean_latent_indices.to(transformer.device),
+                clean_latents_2x=clean_latents_2x.to(
+                    transformer.device, dtype=transformer.dtype
+                ),
+                clean_latent_2x_indices=clean_latent_2x_indices.to(transformer.device),
+                clean_latents_4x=clean_latents_4x.to(
+                    transformer.device, dtype=transformer.dtype
+                ),
+                clean_latent_4x_indices=clean_latent_4x_indices.to(transformer.device),
+                callback=callback_diffusion_step,
+            )
+
+            if is_last_section:
+                generated_latents = torch.cat(
+                    [start_latent.to(generated_latents), generated_latents], dim=2
+                )
+
+            total_generated_latent_frames += int(generated_latents.shape[2])
+            history_latents = torch.cat(
+                [generated_latents.to(history_latents), history_latents], dim=2
+            )
+
+            if not high_vram:
+                offload_model_from_device_for_memory_preservation(
+                    transformer, target_device=gpu, preserved_memory_gb=8
+                )
+                load_model_as_complete(vae, target_device=gpu)
+
+            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 = (
+                    (latent_window_size * 2 + 1)
+                    if is_last_section
+                    else (latent_window_size * 2)
+                )
+                overlapped_frames = latent_window_size * 4 - 3
+                current_pixels = vae_decode(
+                    real_history_latents[:, :, :section_latent_frames], vae
+                ).cpu()
+                history_pixels = soft_append_bcthw(
+                    current_pixels, history_pixels, overlapped_frames
+                )
+
+            if not high_vram:
+                unload_complete_models()
+
+            current_video_frame_count = history_pixels.shape[2]
+
+            # --- Gemini start again
+            # # Skip writing preview mp4 for this segment logic
+            # should_save_mp4_this_iteration = False
+            # current_segment_1_indexed = latent_padding_iteration # + 1
+            # if (latent_padding_iteration == 0) or is_last_section or (parsed_segments_to_decode_set and current_segment_1_indexed in parsed_segments_to_decode_set):
+            #     should_save_mp4_this_iteration = True
+            # if should_save_mp4_this_iteration:
+            #     segment_mp4_filename = os.path.join(outputs_folder, f'{job_id}_segment_{latent_padding_iteration}_frames_{current_video_frame_count}.mp4')
+            #     save_bcthw_as_mp4(history_pixels, segment_mp4_filename, fps=30, crf=mp4_crf)
+            #     final_output_filename = segment_mp4_filename
+            #     print(f"Task {task_id}: SAVED MP4 for segment {latent_padding_iteration} to {segment_mp4_filename}. Total video frames: {current_video_frame_count}")
+            #     output_queue_ref.push(('file', (task_id, segment_mp4_filename, f"Segment {latent_padding_iteration} MP4 saved ({current_video_frame_count} frames)")))
+            # else:
+            #     print(f"Task {task_id}: SKIPPED MP4 save for intermediate segment {current_segment_1_indexed}.")
+
+            # if is_last_section: success = True; break
+
+            # --- Gemini start again
+
+            # ^ original code | v Flash code
+
+            # # Skip writing preview mp4 for this segment logic
+            # should_save_mp4_this_iteration = False
+            # # Use latent_padding_iteration directly here, as it's the 0-indexed loop counter
+            # current_segment_index = latent_padding_iteration
+
+            # # Condition 1: Always save the first segment (index 0)
+            # if current_segment_index == 0:
+            #     should_save_mp4_this_iteration = True
+            # # Condition 2: Always save the last segment
+            # elif is_last_section:
+            #     should_save_mp4_this_iteration = True
+            # # Condition 3: Save if the current segment index is in the parsed set
+            # elif parsed_segments_to_decode_set and (current_segment_index + 1) in parsed_segments_to_decode_set:
+            #     # Add 1 here if segments_to_decode_csv assumes 1-based indexing for user input
+            #     should_save_mp4_this_iteration = True
+            # # Condition 4: Save based on preview_frequency, if enabled (preview_frequency > 0)
+            # elif preview_frequency > 0 and current_segment_index % preview_frequency == 0:
+            #     should_save_mp4_this_iteration = True
+
+            # if should_save_mp4_this_iteration:
+            #     segment_mp4_filename = os.path.join(outputs_folder, f'{job_id}_segment_{latent_padding_iteration}_frames_{current_video_frame_count}.mp4')
+            #     save_bcthw_as_mp4(history_pixels, segment_mp4_filename, fps=30, crf=mp4_crf)
+            #     final_output_filename = segment_mp4_filename
+            #     print(f"Task {task_id}: SAVED MP4 for segment {latent_padding_iteration} to {segment_mp4_filename}. Total video frames: {current_video_frame_count}")
+            #     output_queue_ref.push(('file', (task_id, segment_mp4_filename, f"Segment {latent_padding_iteration} MP4 saved ({current_video_frame_count} frames)")))
+            # else:
+            #     print(f"Task {task_id}: SKIPPED MP4 save for intermediate segment {current_segment_index}.")
+
+            # Determine if we should save an intermediate MP4 for this loop segment
+            should_save_mp4_this_iteration = False
+
+            # Condition 1: Always save the last segment of the loop
+            if is_last_section:
+                should_save_mp4_this_iteration = True
+            # Condition 2: Save if the current loop segment number is explicitly in the parsed set
+            elif (
+                parsed_segments_to_decode_set
+                and current_loop_segment_number in parsed_segments_to_decode_set
+            ):
+                should_save_mp4_this_iteration = True
+            # Condition 3: Save based on preview_frequency, if enabled (preview_frequency > 0)
+            elif preview_frequency > 0 and (
+                current_loop_segment_number % preview_frequency == 0
+            ):
+                should_save_mp4_this_iteration = True
+
+            if should_save_mp4_this_iteration:
+                segment_mp4_filename = os.path.join(
+                    outputs_folder,
+                    f"{job_id}_segment_{current_loop_segment_number}_frames_{current_video_frame_count}.mp4",
+                )  # Updated filename to use 1-indexed segment
+                save_bcthw_as_mp4(
+                    history_pixels, segment_mp4_filename, fps=30, crf=mp4_crf
+                )
+                final_output_filename = segment_mp4_filename
+                print(
+                    f"Task {task_id}: SAVED MP4 for segment {current_loop_segment_number} to {segment_mp4_filename}. Total video frames: {current_video_frame_count}"
+                )  # Updated log to use 1-indexed segment
+                output_queue_ref.push(
+                    (
+                        "file",
+                        (
+                            task_id,
+                            segment_mp4_filename,
+                            f"Segment {current_loop_segment_number} MP4 saved ({current_video_frame_count} frames)",
+                        ),
+                    )
+                )  # Updated output queue message to use 1-indexed segment
+            else:
+                print(
+                    f"Task {task_id}: SKIPPED MP4 save for intermediate segment {current_loop_segment_number}."
+                )  # Updated log to use 1-indexed segment
+
+    except KeyboardInterrupt:
+        print(f"Worker task {task_id} caught KeyboardInterrupt (likely abort signal).")
+        output_queue_ref.push(("aborted", task_id))
+        success = False
+    except Exception as e:
+        print(f"Error in worker task {task_id}: {e}")
+        traceback.print_exc()
+        output_queue_ref.push(("error", (task_id, str(e))))
+        success = False
+    finally:
+        transformer.high_quality_fp32_output_for_inference = original_fp32_setting
+        print(
+            f"Task {task_id}: Restored transformer.high_quality_fp32_output_for_inference to {original_fp32_setting}"
+        )
+        if not high_vram:
+            unload_complete_models(
+                text_encoder, text_encoder_2, image_encoder, vae, transformer
+            )
+        if final_output_filename and not os.path.dirname(
+            final_output_filename
+        ) == os.path.abspath(outputs_folder):
+            final_output_filename = os.path.join(
+                outputs_folder, os.path.basename(final_output_filename)
+            )
+        output_queue_ref.push(("end", (task_id, success, final_output_filename)))

goan.py

@@ -0,0 +1,195 @@
+# goan.py
+# Main application entry point for the goan video generation UI.
+
+# --- Python Standard Library Imports ---
+import os
+import gradio as gr
+import torch
+import argparse
+import atexit
+
+# --- Local Application Imports ---
+# Import managers for different UI sections and shared state
+from ui import layout as layout_manager
+from ui import metadata as metadata_manager
+from ui import queue as queue_manager # Renamed for clarity
+from ui import workspace as workspace_manager
+from ui import shared_state # Ensure this is used consistently
+
+# --- Diffusers and Helper Imports ---
+from diffusers import AutoencoderKLHunyuanVideo
+from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer, SiglipImageProcessor, SiglipVisionModel
+from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
+from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
+from diffusers_helper.gradio.progress_bar import make_progress_bar_css
+
+# --- Environment Setup ---
+os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
+
+# --- Argument Parsing ---
+parser = argparse.ArgumentParser(description="goan: FramePack-based Video Generation UI")
+parser.add_argument('--share', action='store_true', default=False, help="Enable Gradio sharing link.")
+parser.add_argument("--server", type=str, default='127.0.0.1', help="Server name to bind to.")
+parser.add_argument("--port", type=int, required=False, help="Port to run the server on.")
+parser.add_argument("--inbrowser", action='store_true', default=False, help="Launch in browser automatically.")
+# Add the allowed_output_paths argument here
+parser.add_argument("--allowed_output_paths", type=str, default="", help="Comma-separated list of additional output folders Gradio is allowed to access. E.g., '~/my_outputs, /mnt/external_drive/vids'")
+args = parser.parse_args()
+print(f"goan launching with args: {args}")
+
+# --- Model Loading ---
+print("Initializing models...")
+free_mem_gb = get_cuda_free_memory_gb(gpu)
+high_vram = free_mem_gb > 60
+print(f'Free VRAM {free_mem_gb} GB, High-VRAM Mode: {high_vram}')
+
+# Populate shared_state.models with loaded model instances
+shared_state.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(),
+    'transformer': HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu(),
+    'high_vram': high_vram # Renamed key to match worker's expected param
+}
+print("Models loaded to CPU. Configuring...")
+for model_name in ['vae', 'text_encoder', 'text_encoder_2', 'image_encoder', 'transformer']:
+    shared_state.models[model_name].eval()
+if not high_vram:
+    shared_state.models['vae'].enable_slicing(); shared_state.models['vae'].enable_tiling()
+shared_state.models['transformer'].high_quality_fp32_output_for_inference = False
+for model_name, dtype in [('transformer', torch.bfloat16), ('vae', torch.float16), ('image_encoder', torch.float16), ('text_encoder', torch.float16), ('text_encoder_2', torch.float16)]:
+    shared_state.models[model_name].to(dtype=dtype)
+for model_obj in shared_state.models.values(): # Iterate over values, not keys
+    if isinstance(model_obj, torch.nn.Module): model_obj.requires_grad_(False) # Use model_obj here
+if not high_vram:
+    print("Low VRAM mode: Installing DynamicSwap.")
+    DynamicSwapInstaller.install_model(shared_state.models['transformer'], device=gpu)
+    DynamicSwapInstaller.install_model(shared_state.models['text_encoder'], device=gpu)
+else:
+    print("High VRAM mode: Moving all models to GPU.")
+    for model_name in ['text_encoder', 'text_encoder_2', 'image_encoder', 'vae', 'transformer']:
+        shared_state.models[model_name].to(gpu)
+print("Model configuration and placement complete.")
+
+
+# --- UI Helper Functions ---
+def patched_video_is_playable(video_filepath): return True
+gr.processing_utils.video_is_playable = patched_video_is_playable
+
+def ui_update_total_segments(total_seconds_ui, latent_window_size_ui):
+    """Calculates and formats the total segment count for display in the UI."""
+    try:
+        total_segments = int(max(round((total_seconds_ui * 30) / (latent_window_size_ui * 4)), 1))
+        return f"Calculated Total Segments: {total_segments}"
+    except (TypeError, ValueError): return "Segments: Invalid input"
+
+
+# --- UI Creation and Event Wiring ---
+print("Creating UI layout...")
+# Create the UI by calling the layout manager. This returns the block and a dictionary of components.
+ui_components = layout_manager.create_ui()
+block = ui_components['block']
+
+# Define lists of components for easier wiring of events
+creative_ui_keys = ['prompt_ui', 'n_prompt_ui', 'total_second_length_ui', 'seed_ui', 'preview_frequency_ui', 'segments_to_decode_csv_ui', 'gs_ui', 'gs_schedule_shape_ui', 'gs_final_ui', 'steps_ui', 'cfg_ui', 'rs_ui']
+environment_ui_keys = ['use_teacache_ui', 'use_fp32_transformer_output_checkbox_ui', 'gpu_memory_preservation_ui', 'mp4_crf_ui', 'output_folder_ui_ctrl', 'latent_window_size_ui']
+full_workspace_ui_keys = creative_ui_keys + environment_ui_keys
+
+creative_ui_components = [ui_components[key] for key in creative_ui_keys]
+full_workspace_ui_components = [ui_components[key] for key in full_workspace_ui_keys]
+task_defining_ui_inputs = [ui_components['input_image_gallery_ui']] + full_workspace_ui_components
+
+# Define output lists for complex Gradio calls
+process_queue_outputs_list = [ui_components[key] for key in ['app_state', 'queue_df_display_ui', 'last_finished_video_ui', 'current_task_preview_image_ui', 'current_task_progress_desc_ui', 'current_task_progress_bar_ui', 'process_queue_button', 'abort_task_button', 'reset_ui_button']]
+queue_df_select_outputs_list = [ui_components[key] for key in ['app_state', 'queue_df_display_ui', 'input_image_gallery_ui'] + full_workspace_ui_keys + ['add_task_button', 'cancel_edit_task_button', 'last_finished_video_ui']]
+
+# Wire up all the UI events to their handler functions in the respective managers
+with block:
+    # Workspace Manager Events
+    ui_components['save_workspace_button'].click(fn=workspace_manager.save_workspace, inputs=full_workspace_ui_components, outputs=None)
+    ui_components['load_workspace_button'].click(fn=workspace_manager.load_workspace, inputs=None, outputs=full_workspace_ui_components)
+    ui_components['save_as_default_button'].click(fn=workspace_manager.save_as_default_workspace, inputs=full_workspace_ui_components, outputs=None)
+
+    # Metadata Manager Events
+    ui_components['input_image_gallery_ui'].upload(fn=metadata_manager.handle_image_upload_for_metadata, inputs=[ui_components['input_image_gallery_ui']], outputs=[ui_components['metadata_modal']])
+    ui_components['confirm_metadata_btn'].click(fn=metadata_manager.apply_and_hide_modal, inputs=[ui_components['input_image_gallery_ui']], outputs=[ui_components['metadata_modal']] + creative_ui_components)
+    ui_components['cancel_metadata_btn'].click(fn=lambda: gr.update(visible=False), inputs=None, outputs=ui_components['metadata_modal'])
+
+    # Queue Manager Events
+    ui_components['add_task_button'].click(fn=queue_manager.add_or_update_task_in_queue, inputs=[ui_components['app_state']] + task_defining_ui_inputs, outputs=[ui_components['app_state'], ui_components['queue_df_display_ui'], ui_components['add_task_button'], ui_components['cancel_edit_task_button']])
+    ui_components['process_queue_button'].click(fn=queue_manager.process_task_queue_main_loop, inputs=[ui_components['app_state']], outputs=process_queue_outputs_list)
+    ui_components['cancel_edit_task_button'].click(fn=queue_manager.cancel_edit_mode_action, inputs=[ui_components['app_state']], outputs=[ui_components['app_state'], ui_components['queue_df_display_ui'], ui_components['add_task_button'], ui_components['cancel_edit_task_button']])
+    ui_components['abort_task_button'].click(fn=queue_manager.abort_current_task_processing_action, inputs=[ui_components['app_state']], outputs=[ui_components['app_state'], ui_components['abort_task_button']])
+    ui_components['clear_queue_button_ui'].click(fn=queue_manager.clear_task_queue_action, inputs=[ui_components['app_state']], outputs=[ui_components['app_state'], ui_components['queue_df_display_ui']])
+    ui_components['save_queue_button_ui'].click(fn=queue_manager.save_queue_to_zip, inputs=[ui_components['app_state']], outputs=[ui_components['app_state'], ui_components['save_queue_zip_b64_output']]).then(fn=None, inputs=[ui_components['save_queue_zip_b64_output']], outputs=None, js="""(b64) => { if(!b64) return; const blob = new Blob([Uint8Array.from(atob(b64), c => c.charCodeAt(0))], {type: 'application/zip'}); const url = URL.createObjectURL(blob); const a = document.createElement('a'); a.href=url; a.download='goan_queue.zip'; a.click(); URL.revokeObjectURL(url); }""")
+    ui_components['load_queue_button_ui'].upload(fn=queue_manager.load_queue_from_zip, inputs=[ui_components['app_state'], ui_components['load_queue_button_ui']], outputs=[ui_components['app_state'], ui_components['queue_df_display_ui']])
+    ui_components['queue_df_display_ui'].select(fn=queue_manager.handle_queue_action_on_select, inputs=[ui_components['app_state']] + task_defining_ui_inputs, outputs=queue_df_select_outputs_list)
+
+    # Other UI Event Handlers
+    ui_components['gs_schedule_shape_ui'].change(fn=lambda choice: gr.update(interactive=(choice != "Off")), inputs=[ui_components['gs_schedule_shape_ui']], outputs=[ui_components['gs_final_ui']])
+    for ctrl_key in ['total_second_length_ui', 'latent_window_size_ui']:
+        ui_components[ctrl_key].change(fn=ui_update_total_segments, inputs=[ui_components['total_second_length_ui'], ui_components['latent_window_size_ui']], outputs=[ui_components['total_segments_display_ui']])
+
+    refresh_image_path_state = gr.State(None)
+    # The reset_ui_button's functionality remains the same: save state then reload page
+    ui_components['reset_ui_button'].click(fn=workspace_manager.save_ui_and_image_for_refresh, inputs=task_defining_ui_inputs, outputs=None).then(fn=None, inputs=None, outputs=None, js="() => { window.location.reload(); }")
+
+    # --- Application Startup and Shutdown ---
+    autoload_outputs = [ui_components[k] for k in ['app_state', 'queue_df_display_ui', 'process_queue_button', 'abort_task_button', 'last_finished_video_ui']]
+
+    # This is the crucial block.load chain to ensure re-attachment
+    (block.load(fn=workspace_manager.load_workspace_on_start, inputs=[], outputs=[refresh_image_path_state] + full_workspace_ui_components)
+        .then(fn=workspace_manager.load_image_from_path, inputs=[refresh_image_path_state], outputs=[ui_components['input_image_gallery_ui']])
+        .then(fn=queue_manager.autoload_queue_on_start_action, inputs=[ui_components['app_state']], outputs=autoload_outputs)
+        .then(lambda s_val: shared_state.global_state_for_autosave.update(s_val), inputs=[ui_components['app_state']], outputs=None)
+        .then(fn=ui_update_total_segments, inputs=[ui_components['total_second_length_ui'], ui_components['latent_window_size_ui']], outputs=[ui_components['total_segments_display_ui']])
+        # *** ADDED: Automatic re-attachment of progress UI on load if processing ***
+        .then(
+            fn=queue_manager.process_task_queue_main_loop,
+            inputs=[ui_components['app_state']],
+            outputs=process_queue_outputs_list, # Use the existing output list
+            js="""
+            (app_state_val) => {
+                // This JS runs after autoload_queue_on_start_action completes.
+                // If a task is processing, we want to re-invoke the Python generator.
+                if (app_state_val.queue_state && app_state_val.queue_state.processing) {
+                    console.log("Gradio: Auto-reconnecting to ongoing task output stream.");
+                    // Return a non-null, non-falsey value to trigger the Python function.
+                    return "reconnect_stream";
+                }
+                console.log("Gradio: No ongoing task detected for auto-reconnection.");
+                return null; // Return null to skip calling the Python function
+            }
+            """
+        )
+    )
+
+    # Register the atexit handler with the global_state_for_autosave
+    atexit.register(queue_manager.autosave_queue_on_exit_action, shared_state.global_state_for_autosave)
+
+# --- Application Launch ---
+if __name__ == "__main__":
+    print("Starting goan FramePack UI...")
+    # Determine the initial output folder for allowed_paths based on saved settings or default
+    initial_output_folder_path = workspace_manager.get_initial_output_folder_from_settings()
+    expanded_outputs_folder_for_launch = os.path.abspath(initial_output_folder_path)
+
+    # Prepare the list of allowed paths for Gradio
+    final_allowed_paths = [expanded_outputs_folder_for_launch]
+
+    if args.allowed_output_paths:
+        custom_cli_paths = [
+            os.path.abspath(os.path.expanduser(p.strip()))
+            for p in args.allowed_output_paths.split(',')
+            if p.strip()
+        ]
+        final_allowed_paths.extend(custom_cli_paths)
+
+    final_allowed_paths = list(set(final_allowed_paths)) # Remove duplicates
+
+    print(f"Gradio allowed paths: {final_allowed_paths}")
+    block.launch(server_name=args.server, server_port=args.port, share=args.share, inbrowser=args.inbrowser, allowed_paths=final_allowed_paths)

requirements.txt

@@ -0,0 +1,20 @@
+accelerate==1.6.0
+diffusers==0.33.1
+transformers==4.46.2
+gradio==5.23.0
+gradio-modal
+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
+
+einops
+opencv-contrib-python
+safetensors
+
+# --- OK if one fails but less so if they both do ---
+xformers
+sageattention

ui/__init__.py

@@ -0,0 +1 @@
+ 

ui/layout.py

@@ -0,0 +1,126 @@
+# ui/layout.py
+# This file defines the Gradio UI layout for the goan application.
+
+import gradio as gr
+from gradio_modal import Modal
+
+# Import workspace manager to get the default output folder path
+from . import workspace as workspace_manager
+
+def create_ui():
+    """
+    Creates the Gradio UI layout and returns a dictionary of all UI components.
+    This separation of layout from logic makes the main script cleaner.
+    """
+    # Define a dictionary to hold all UI components, which will be returned.
+    components = {}
+
+    css = """
+    #queue_df { font-size: 0.85rem; } 
+    #queue_df th:nth-child(1), #queue_df td:nth-child(1) { width: 5%; } 
+    #queue_df th:nth-child(2), #queue_df td:nth-child(2) { width: 10%; } 
+    #queue_df th:nth-child(3), #queue_df td:nth-child(3) { width: 40%; overflow: hidden; text-overflow: ellipsis; white-space: nowrap;} 
+    #queue_df th:nth-child(4), #queue_df td:nth-child(4) { width: 8%; } 
+    #queue_df th:nth-child(5), #queue_df td:nth-child(5) { width: 8%; } 
+    #queue_df th:nth-child(6), #queue_df td:nth-child(6) { width: 10%; text-align: center; } 
+    #queue_df th:nth-child(7), #queue_df td:nth-child(7), 
+    #queue_df th:nth-child(8), #queue_df td:nth-child(8), 
+    #queue_df th:nth-child(9), #queue_df td:nth-child(9), 
+    #queue_df th:nth-child(10), #queue_df td:nth-child(10) { width: 4%; cursor: pointer; text-align: center; } 
+    #queue_df td:hover { background-color: #f0f0f0; } 
+    .gradio-container { max-width: 95% !important; margin: auto !important; }
+    """
+    
+    block = gr.Blocks(css=css, title="goan").queue()
+    components['block'] = block
+
+    with block:
+        # The app_state dictionary holds all transient state for the UI
+        app_state = gr.State({
+            "queue_state": {"queue": [], "next_id": 1, "processing": False, "editing_task_id": None},
+            "last_completed_video_path": None
+        })
+        components['app_state'] = app_state
+
+        gr.Markdown('# goan (Powered by FramePack)')
+
+        with Modal(visible=False) as metadata_modal:
+            components['metadata_modal'] = metadata_modal
+            gr.Markdown("Image has saved parameters. Overwrite current creative settings?")
+            with gr.Row():
+                components['cancel_metadata_btn'] = gr.Button("No")
+                components['confirm_metadata_btn'] = gr.Button("Yes, Apply", variant="primary")
+
+        # --- Start of layout structure ---
+
+        # TOP ZONE: A row with two columns.
+        # Column 1 (scale=1): Image input and primary queue actions.
+        # Column 2 (scale=2): Prompting.
+        with gr.Row():
+            # Column 1: Input Image & Primary Actions
+            with gr.Column(scale=1, min_width=300):
+                components['input_image_gallery_ui'] = gr.Gallery(type="pil", label="Input Image(s)", height=220)
+                components['add_task_button'] = gr.Button("Add to Queue", variant="secondary")
+                components['process_queue_button'] = gr.Button("▶️ Process Queue", variant="primary")
+                components['abort_task_button'] = gr.Button("⏹️ Abort", variant="stop", interactive=True) # default to active on startup in case service is still running
+                components['cancel_edit_task_button'] = gr.Button("Cancel Edit", visible=False, variant="secondary")
+
+            # Column 2: Prompts
+            with gr.Column(scale=2, min_width=600):
+                components['prompt_ui'] = gr.Textbox(label="Prompt", lines=10)
+                components['n_prompt_ui'] = gr.Textbox(label="Negative Prompt", lines=4)
+
+        # MIDDLE ZONE: Full-width task queue display and file operations
+        with gr.Group():
+            gr.Markdown("## Task Queue")
+            components['queue_df_display_ui'] = gr.DataFrame(headers=["ID", "Status", "Prompt", "Length", "Steps", "Input", "↑", "↓", "✖", "✎"], datatype=["number","markdown","markdown","str","number","markdown","markdown","markdown","markdown","markdown"], col_count=(10,"fixed"), interactive=False, visible=False, elem_id="queue_df")
+            with gr.Row():
+                components['save_queue_zip_b64_output'] = gr.Text(visible=False)
+                components['save_queue_button_ui'] = gr.DownloadButton("Save Queue", size="sm")
+                components['load_queue_button_ui'] = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm")
+                components['clear_queue_button_ui'] = gr.Button("Clear Pending", size="sm", variant="stop")
+
+        # BOTTOM ZONE: A row split into two columns for Settings and Output
+        with gr.Row(equal_height=False):
+            # Column 1: All settings and parameters
+            with gr.Column(scale=1):
+                with gr.Row():
+                    components['total_second_length_ui'] = gr.Slider(label="Video Length (s)", minimum=0.1, maximum=120, value=5.0, step=0.1)
+                    components['seed_ui'] = gr.Number(label="Seed", value=-1, precision=0)
+                with gr.Accordion("Advanced Settings", open=False): # keep default closed on startup
+                    components['total_segments_display_ui'] = gr.Markdown("Calculated Total Segments: N/A")
+                    components['preview_frequency_ui'] = gr.Slider(label="Preview Freq.", minimum=0, maximum=100, value=5, step=1)
+                    components['segments_to_decode_csv_ui'] = gr.Textbox(label="Preview Segments CSV", value="")
+                    with gr.Row():
+                        components['gs_ui'] = gr.Slider(label="Distilled CFG Start", minimum=1.0, maximum=32.0, value=10.0, step=0.01)
+                        components['gs_schedule_shape_ui'] = gr.Radio(["Off", "Linear"], label="Variable CFG", value="Off")
+                        components['gs_final_ui'] = gr.Slider(label="Distilled CFG End", minimum=1.0, maximum=32.0, value=10.0, step=0.01, interactive=False)
+                    components['cfg_ui'] = gr.Slider(label="CFG (Real)", minimum=1.0, maximum=32.0, value=1.0, step=0.01)
+                    components['steps_ui'] = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)
+                    components['rs_ui'] = gr.Slider(label="RS", minimum=0.0, maximum=32.0, value=0.0, step=0.01, visible=False)
+                with gr.Accordion("Debug Settings", open=False): # keep default closed on startup
+                    components['use_teacache_ui'] = gr.Checkbox(label='Use TeaCache', value=True)
+                    components['use_fp32_transformer_output_checkbox_ui'] = gr.Checkbox(label="Use FP32 Transformer Output", value=False)
+                    components['gpu_memory_preservation_ui'] = gr.Slider(label="GPU Preserved (GB)", minimum=4, maximum=128, value=6.0, step=0.1)
+                    components['mp4_crf_ui'] = gr.Slider(label="MP4 CRF", minimum=0, maximum=51, value=18, step=1)
+                    components['latent_window_size_ui'] = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False)
+                    components['output_folder_ui_ctrl'] = gr.Textbox(label="Output Folder", value=workspace_manager.outputs_folder)
+                    components['save_as_default_button'] = gr.Button("Save as Default", variant="secondary")
+                    components['reset_ui_button'] = gr.Button("Save & Refresh UI", variant="secondary")
+
+                
+                # UI and Workspace Management Buttons
+                with gr.Row():
+                    components['save_workspace_button'] = gr.Button("Save Workspace", variant="secondary")
+                    components['load_workspace_button'] = gr.Button("Load Workspace", variant="secondary")
+
+            # Column 2: Live Preview and Final Output
+            with gr.Column(scale=1):
+                gr.Markdown("## Live Preview & Output")
+                components['current_task_preview_image_ui'] = gr.Image(interactive=False, visible=False)
+                components['current_task_progress_desc_ui'] = gr.Markdown('')
+                components['current_task_progress_bar_ui'] = gr.HTML('')
+                # Set a height on the video player to help with layout stability
+                components['last_finished_video_ui'] = gr.Video(interactive=True, autoplay=False, height=540)
+                
+    return components

ui/metadata.py

@@ -0,0 +1,88 @@
+# ui/metadata.py
+import json
+import gradio as gr
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+
+# Import shared state for access to parameter key lists
+from . import shared_state
+
+# --- Core Metadata Functions ---
+
+def extract_metadata_from_pil_image(pil_image: Image.Image) -> dict:
+    """Extracts a 'parameters' dictionary from a PIL image's text chunk."""
+    if pil_image is None: 
+        return {}
+    
+    pnginfo_data = getattr(pil_image, 'text', None)
+    if not isinstance(pnginfo_data, dict):
+        return {}
+    
+    params_json_str = pnginfo_data.get('parameters')
+    if not params_json_str:
+        return {}
+    
+    try:
+        extracted_params = json.loads(params_json_str)
+        return extracted_params if isinstance(extracted_params, dict) else {}
+    except json.JSONDecodeError as e:
+        print(f"Error decoding metadata JSON: {e}")
+        return {}
+
+def write_image_metadata(pil_image: Image.Image, params_dict: dict) -> Image.Image:
+    """Creates a PngInfo object with the given parameters and attaches it to the image."""
+    metadata = PngInfo()
+    metadata.add_text("parameters", json.dumps(params_dict))
+    pil_image.info = metadata # The .info attribute is the correct place to assign the PngInfo object
+    return pil_image
+
+# --- UI Handler Functions (Moved from demo_gradio_svc.py) ---
+
+def handle_image_upload_for_metadata(gallery_pil_list):
+    """
+    Checks an uploaded image for metadata and shows a confirmation modal if found.
+    This function is triggered by the 'upload' event of the image gallery.
+    """
+    if not gallery_pil_list:
+        return gr.update(visible=False)
+    
+    # The gallery component returns a list of (image, name) tuples.
+    pil_image = gallery_pil_list[0][0] if isinstance(gallery_pil_list[0], tuple) else gallery_pil_list[0]
+    
+    if isinstance(pil_image, Image.Image):
+        extracted_metadata = extract_metadata_from_pil_image(pil_image)
+        # Show the modal only if metadata exists and contains relevant keys.
+        if extracted_metadata and any(key in extracted_metadata for key in shared_state.CREATIVE_PARAM_KEYS):
+            return gr.update(visible=True)
+            
+    return gr.update(visible=False)
+
+def ui_load_params_from_image_metadata(gallery_data_list):
+    """
+    Loads ONLY the creative parameters from image metadata and returns UI updates.
+    """
+    updates = [gr.update()] * len(shared_state.CREATIVE_PARAM_KEYS)
+    if not gallery_data_list:
+        return updates
+        
+    pil_image = gallery_data_list[0][0] if isinstance(gallery_data_list[0], tuple) else gallery_data_list[0]
+    extracted_metadata = extract_metadata_from_pil_image(pil_image)
+    
+    if not extracted_metadata:
+        gr.Info("No parameters found in image.")
+        return updates
+        
+    gr.Info(f"Applying creative settings from image...")
+    for i, key in enumerate(shared_state.CREATIVE_PARAM_KEYS):
+        if key in extracted_metadata:
+            # Return a Gradio update object for each changed parameter.
+            updates[i] = gr.update(value=extracted_metadata[key])
+            
+    return updates
+
+def apply_and_hide_modal(gallery_data_list):
+    """
+    A wrapper function that applies the metadata and then hides the confirmation modal.
+    """
+    # The first output hides the modal, the rest are the parameter updates.
+    return [gr.update(visible=False)] + ui_load_params_from_image_metadata(gallery_data_list)

ui/queue.py

@@ -0,0 +1,397 @@
+# ui/queue.py
+import gradio as gr
+import numpy as np
+from PIL import Image
+import os
+import json
+import base64
+import io
+import zipfile
+import tempfile
+import atexit
+import traceback
+from pathlib import Path
+
+# Import shared state and constants from the dedicated module.
+from . import shared_state
+from generation_core import worker
+from diffusers_helper.thread_utils import AsyncStream, async_run
+
+
+# Configuration for the autosave feature.
+AUTOSAVE_FILENAME = "goan_autosave_queue.zip"
+
+
+def np_to_base64_uri(np_array_or_tuple, format="png"):
+    """Converts a NumPy array representing an image to a base64 data URI."""
+    if np_array_or_tuple is None: return None
+    try:
+        np_array = np_array_or_tuple[0] if isinstance(np_array_or_tuple, tuple) and len(np_array_or_tuple) > 0 and isinstance(np_array_or_tuple[0], np.ndarray) else np_array_or_tuple if isinstance(np_array_or_tuple, np.ndarray) else None
+        if np_array is None: return None
+        pil_image = Image.fromarray(np_array.astype(np.uint8))
+        if format.lower() == "jpeg" and pil_image.mode == "RGBA": pil_image = pil_image.convert("RGB")
+        buffer = io.BytesIO(); pil_image.save(buffer, format=format.upper()); img_bytes = buffer.getvalue()
+        return f"data:image/{format.lower()};base64,{base64.b64encode(img_bytes).decode('utf-8')}"
+    except Exception as e: print(f"Error converting NumPy to base64: {e}"); return None
+
+def get_queue_state(state_dict_gr_state):
+    """Safely retrieves the queue_state dictionary from the main application state."""
+    if "queue_state" not in state_dict_gr_state: state_dict_gr_state["queue_state"] = {"queue": [], "next_id": 1, "processing": False, "editing_task_id": None}
+    return state_dict_gr_state["queue_state"]
+
+def update_queue_df_display(queue_state):
+    """Generates a Gradio DataFrame update from the current queue state."""
+    queue = queue_state.get("queue", []); data = []; processing = queue_state.get("processing", False); editing_task_id = queue_state.get("editing_task_id", None)
+    for i, task in enumerate(queue):
+        params = task['params']; task_id = task['id']; prompt_display = (params['prompt'][:77] + '...') if len(params['prompt']) > 80 else params['prompt']; prompt_title = params['prompt'].replace('"', '"'); prompt_cell = f'<span title="{prompt_title}">{prompt_display}</span>'; img_uri = np_to_base64_uri(params.get('input_image'), format="png"); thumbnail_size = "50px"; img_md = f'<img src="{img_uri}" alt="Input" style="max-width:{thumbnail_size}; max-height:{thumbnail_size}; display:block; margin:auto; object-fit:contain;" />' if img_uri else ""; is_processing_current_task = processing and i == 0; is_editing_current_task = editing_task_id == task_id; task_status_val = task.get("status", "pending");
+        if is_processing_current_task: status_display = "⏳ Processing"
+        elif is_editing_current_task: status_display = "✏️ Editing"
+        elif task_status_val == "done": status_display = "✅ Done"
+        elif task_status_val == "error": status_display = f"❌ Error: {task.get('error_message', 'Unknown')}"
+        elif task_status_val == "aborted": status_display = "⏹️ Aborted"
+        elif task_status_val == "pending": status_display = "⏸️ Pending"
+        data.append([task_id, status_display, prompt_cell, f"{params.get('total_second_length', 0):.1f}s", params.get('steps', 0), img_md, "↑", "↓", "✖", "✎"])
+    return gr.DataFrame(value=data, visible=len(data) > 0)
+
+def add_or_update_task_in_queue(state_dict_gr_state, *args_from_ui_controls_tuple):
+    """Adds a new task to the queue or updates an existing one if in edit mode."""
+    queue_state = get_queue_state(state_dict_gr_state); editing_task_id = queue_state.get("editing_task_id", None)
+
+    input_images_pil_list = args_from_ui_controls_tuple[0]
+    all_ui_values_tuple = args_from_ui_controls_tuple[1:]
+    if not input_images_pil_list:
+        gr.Warning("Input image is required!")
+        return state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value="Add Task to Queue" if editing_task_id is None else "Update Task"), gr.update(visible=editing_task_id is not None)
+
+    temp_params_from_ui = dict(zip(shared_state.ALL_TASK_UI_KEYS, all_ui_values_tuple))
+    base_params_for_worker_dict = {}
+    for ui_key, worker_key in shared_state.UI_TO_WORKER_PARAM_MAP.items():
+        if ui_key == 'gs_schedule_shape_ui':
+            base_params_for_worker_dict[worker_key] = temp_params_from_ui.get(ui_key) != 'Off'
+        else:
+            base_params_for_worker_dict[worker_key] = temp_params_from_ui.get(ui_key)
+
+    if editing_task_id is not None:
+        if len(input_images_pil_list) > 1: gr.Warning("Cannot update task with multiple images. Cancel edit."); return state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value="Update Task"), gr.update(visible=True)
+        pil_img_for_update = input_images_pil_list[0][0] if isinstance(input_images_pil_list[0], tuple) else input_images_pil_list[0]
+        if not isinstance(pil_img_for_update, Image.Image): gr.Warning("Invalid image format for update."); return state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value="Update Task"), gr.update(visible=True)
+        img_np_for_update = np.array(pil_img_for_update)
+        with shared_state.queue_lock:
+            task_found = False
+            for task in queue_state["queue"]:
+                if task["id"] == editing_task_id:
+                    task["params"] = {**base_params_for_worker_dict, 'input_image': img_np_for_update}
+                    task["status"] = "pending"
+                    task_found = True
+                    break
+            if not task_found: gr.Warning(f"Task {editing_task_id} not found for update.")
+            else: gr.Info(f"Task {editing_task_id} updated.")
+            queue_state["editing_task_id"] = None
+    else:
+        tasks_added_count = 0; first_new_task_id = -1
+        with shared_state.queue_lock:
+            for img_obj in input_images_pil_list:
+                pil_image = img_obj[0] if isinstance(img_obj, tuple) else img_obj
+                if not isinstance(pil_image, Image.Image): gr.Warning("Skipping invalid image input."); continue
+                img_np_data = np.array(pil_image)
+                next_id = queue_state["next_id"]
+                if first_new_task_id == -1: first_new_task_id = next_id
+                task = {"id": next_id, "params": {**base_params_for_worker_dict, 'input_image': img_np_data}, "status": "pending"}
+                queue_state["queue"].append(task); queue_state["next_id"] += 1; tasks_added_count += 1
+        if tasks_added_count > 0: gr.Info(f"Added {tasks_added_count} task(s) (start ID: {first_new_task_id}).")
+        else: gr.Warning("No valid tasks added.")
+
+    return state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value="Add Task(s) to Queue", variant="secondary"), gr.update(visible=False)
+
+def cancel_edit_mode_action(state_dict_gr_state):
+    """Cancels the current task editing session."""
+    queue_state = get_queue_state(state_dict_gr_state)
+    if queue_state.get("editing_task_id") is not None: gr.Info("Edit cancelled."); queue_state["editing_task_id"] = None
+    return state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value="Add Task(s) to Queue", variant="secondary"), gr.update(visible=False)
+
+def move_task_in_queue(state_dict_gr_state, direction: str, selected_indices_list: list):
+    """Moves a selected task up or down in the queue."""
+    if not selected_indices_list or not selected_indices_list[0]: return state_dict_gr_state, update_queue_df_display(get_queue_state(state_dict_gr_state))
+    idx = int(selected_indices_list[0][0]); queue_state = get_queue_state(state_dict_gr_state); queue = queue_state["queue"]
+    with shared_state.queue_lock:
+        if direction == 'up' and idx > 0: queue[idx], queue[idx-1] = queue[idx-1], queue[idx]
+        elif direction == 'down' and idx < len(queue) - 1: queue[idx], queue[idx+1] = queue[idx+1], queue[idx]
+    return state_dict_gr_state, update_queue_df_display(queue_state)
+
+def remove_task_from_queue(state_dict_gr_state, selected_indices_list: list):
+    """Removes a selected task from the queue."""
+    removed_task_id = None
+    if not selected_indices_list or not selected_indices_list[0]: return state_dict_gr_state, update_queue_df_display(get_queue_state(state_dict_gr_state)), removed_task_id
+    idx = int(selected_indices_list[0][0]); queue_state = get_queue_state(state_dict_gr_state); queue = queue_state["queue"]
+    with shared_state.queue_lock:
+        if 0 <= idx < len(queue): removed_task = queue.pop(idx); removed_task_id = removed_task['id']; gr.Info(f"Removed task {removed_task_id}.")
+        else: gr.Warning("Invalid index for removal.")
+    return state_dict_gr_state, update_queue_df_display(queue_state), removed_task_id
+
+def handle_queue_action_on_select(evt: gr.SelectData, state_dict_gr_state, *ui_param_controls_tuple):
+    """Handles clicks on the action buttons (↑, ↓, ✖, ✎) in the queue DataFrame."""
+    if evt.index is None or evt.value not in ["↑", "↓", "✖", "✎"]:
+        return [state_dict_gr_state, update_queue_df_display(get_queue_state(state_dict_gr_state))] + [gr.update()] * (len(shared_state.ALL_TASK_UI_KEYS) + 4)
+
+    row_index, col_index = evt.index; button_clicked = evt.value; queue_state = get_queue_state(state_dict_gr_state); queue = queue_state["queue"]; processing = queue_state.get("processing", False)
+    outputs_list = [state_dict_gr_state, update_queue_df_display(queue_state)] + [gr.update()] * (len(shared_state.ALL_TASK_UI_KEYS) + 4)
+
+    if button_clicked == "↑":
+        if processing and row_index == 0: gr.Warning("Cannot move processing task."); return outputs_list
+        new_state, new_df = move_task_in_queue(state_dict_gr_state, 'up', [[row_index, col_index]]); outputs_list[0], outputs_list[1] = new_state, new_df
+    elif button_clicked == "↓":
+        if processing and row_index == 0: gr.Warning("Cannot move processing task."); return outputs_list
+        if processing and row_index == 1: gr.Warning("Cannot move below processing task."); return outputs_list
+        new_state, new_df = move_task_in_queue(state_dict_gr_state, 'down', [[row_index, col_index]]); outputs_list[0], outputs_list[1] = new_state, new_df
+    elif button_clicked == "✖":
+        if processing and row_index == 0: gr.Warning("Cannot remove processing task."); return outputs_list
+        new_state, new_df, removed_id = remove_task_from_queue(state_dict_gr_state, [[row_index, col_index]]); outputs_list[0], outputs_list[1] = new_state, new_df
+        if removed_id is not None and queue_state.get("editing_task_id", None) == removed_id:
+            queue_state["editing_task_id"] = None
+            outputs_list[2 + 1 + len(shared_state.ALL_TASK_UI_KEYS)] = gr.update(value="Add Task(s) to Queue", variant="secondary")
+            outputs_list[2 + 1 + len(shared_state.ALL_TASK_UI_KEYS) + 1] = gr.update(visible=False)
+    elif button_clicked == "✎":
+        if processing and row_index == 0: gr.Warning("Cannot edit processing task."); return outputs_list
+        if 0 <= row_index < len(queue):
+            task_to_edit = queue[row_index]; task_id_to_edit = task_to_edit['id']; params_to_load_to_ui = task_to_edit['params']
+            queue_state["editing_task_id"] = task_id_to_edit; gr.Info(f"Editing Task {task_id_to_edit}.")
+            img_np_from_task = params_to_load_to_ui.get('input_image')
+            outputs_list[2] = gr.update(value=[(Image.fromarray(img_np_from_task), "loaded_image")]) if isinstance(img_np_from_task, np.ndarray) else gr.update(value=None)
+            for i, ui_key in enumerate(shared_state.ALL_TASK_UI_KEYS):
+                worker_key = shared_state.UI_TO_WORKER_PARAM_MAP.get(ui_key)
+                if worker_key in params_to_load_to_ui:
+                    value_from_task = params_to_load_to_ui[worker_key]
+                    outputs_list[3 + i] = gr.update(value="Linear" if value_from_task else "Off") if ui_key == 'gs_schedule_shape_ui' else gr.update(value=value_from_task)
+            outputs_list[2 + 1 + len(shared_state.ALL_TASK_UI_KEYS)] = gr.update(value="Update Task", variant="secondary")
+            outputs_list[2 + 1 + len(shared_state.ALL_TASK_UI_KEYS) + 1] = gr.update(visible=True)
+        else: gr.Warning("Invalid index for edit.")
+    return outputs_list
+
+def clear_task_queue_action(state_dict_gr_state):
+    """Clears all non-processing tasks from the queue."""
+    queue_state = get_queue_state(state_dict_gr_state); queue = queue_state["queue"]; processing = queue_state["processing"]; cleared_count = 0
+    with shared_state.queue_lock:
+        if processing:
+             if len(queue) > 1: cleared_count = len(queue) - 1; queue_state["queue"] = [queue[0]]; gr.Info(f"Cleared {cleared_count} pending tasks.")
+             else: gr.Info("Only processing task in queue.")
+        elif queue: cleared_count = len(queue); queue.clear(); gr.Info(f"Cleared {cleared_count} tasks.")
+        else: gr.Info("Queue empty.")
+    if not processing and cleared_count > 0 and os.path.isfile(AUTOSAVE_FILENAME):
+         try: os.remove(AUTOSAVE_FILENAME); print(f"Cleared autosave: {AUTOSAVE_FILENAME}.")
+         except OSError as e: print(f"Error deleting autosave: {e}")
+    return state_dict_gr_state, update_queue_df_display(queue_state)
+
+def save_queue_to_zip(state_dict_gr_state):
+    """Saves the current task queue to a zip file for download."""
+    queue_state = get_queue_state(state_dict_gr_state); queue = queue_state.get("queue", [])
+    if not queue: gr.Info("Queue is empty. Nothing to save."); return state_dict_gr_state, ""
+    zip_buffer = io.BytesIO(); saved_files_count = 0
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            queue_manifest = []; image_paths_in_zip = {}
+            for task in queue:
+                params_copy = task['params'].copy(); task_id_s = task['id']; input_image_np_data = params_copy.pop('input_image', None)
+                manifest_entry = {"id": task_id_s, "params": params_copy, "status": task.get("status", "pending")}
+                if input_image_np_data is not None:
+                    img_hash = hash(input_image_np_data.tobytes()); img_filename_in_zip = f"task_{task_id_s}_input.png"; manifest_entry['image_ref'] = img_filename_in_zip
+                    if img_hash not in image_paths_in_zip:
+                        img_save_path = os.path.join(tmpdir, img_filename_in_zip)
+                        try: Image.fromarray(input_image_np_data).save(img_save_path, "PNG"); image_paths_in_zip[img_hash] = img_filename_in_zip; saved_files_count +=1
+                        except Exception as e: print(f"Error saving image for task {task_id_s} in zip: {e}")
+                queue_manifest.append(manifest_entry)
+            manifest_path = os.path.join(tmpdir, "queue_manifest.json");
+            with open(manifest_path, 'w', encoding='utf-8') as f: json.dump(queue_manifest, f, indent=4)
+            with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf:
+                zf.write(manifest_path, arcname="queue_manifest.json")
+                for img_hash, img_filename_rel in image_paths_in_zip.items(): zf.write(os.path.join(tmpdir, img_filename_rel), arcname=img_filename_rel)
+            zip_buffer.seek(0); zip_base64 = base64.b64encode(zip_buffer.getvalue()).decode('utf-8')
+            gr.Info(f"Queue with {len(queue)} tasks ({saved_files_count} images) prepared for download.")
+            return state_dict_gr_state, zip_base64
+    except Exception as e: print(f"Error creating zip for queue: {e}"); traceback.print_exc(); gr.Warning("Failed to create zip data."); return state_dict_gr_state, ""
+    finally: zip_buffer.close()
+
+def load_queue_from_zip(state_dict_gr_state, uploaded_zip_file_obj):
+    """Loads a task queue from an uploaded zip file."""
+    if not uploaded_zip_file_obj or not hasattr(uploaded_zip_file_obj, 'name') or not Path(uploaded_zip_file_obj.name).is_file(): gr.Warning("No valid file selected."); return state_dict_gr_state, update_queue_df_display(get_queue_state(state_dict_gr_state))
+    queue_state = get_queue_state(state_dict_gr_state); newly_loaded_queue = []; max_id_in_file = 0; loaded_image_count = 0; error_messages = []
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir_extract:
+            with zipfile.ZipFile(uploaded_zip_file_obj.name, 'r') as zf:
+                if "queue_manifest.json" not in zf.namelist(): raise ValueError("queue_manifest.json not found in zip")
+                zf.extractall(tmpdir_extract)
+            manifest_path = os.path.join(tmpdir_extract, "queue_manifest.json")
+            with open(manifest_path, 'r', encoding='utf-8') as f: loaded_manifest = json.load(f)
+
+            for task_data in loaded_manifest:
+                params_from_manifest = task_data.get('params', {}); task_id_loaded = task_data.get('id', 0); max_id_in_file = max(max_id_in_file, task_id_loaded)
+                image_ref_from_manifest = task_data.get('image_ref'); input_image_np_data = None
+                if image_ref_from_manifest:
+                    img_path_in_extract = os.path.join(tmpdir_extract, image_ref_from_manifest)
+                    if os.path.exists(img_path_in_extract):
+                        try: input_image_np_data = np.array(Image.open(img_path_in_extract)); loaded_image_count +=1
+                        except Exception as img_e: error_messages.append(f"Err loading img for task {task_id_loaded}: {img_e}")
+                    else: error_messages.append(f"Missing img file for task {task_id_loaded}: {image_ref_from_manifest}")
+                runtime_task = {"id": task_id_loaded, "params": {**params_from_manifest, 'input_image': input_image_np_data}, "status": "pending"}
+                newly_loaded_queue.append(runtime_task)
+        with shared_state.queue_lock: queue_state["queue"] = newly_loaded_queue; queue_state["next_id"] = max(max_id_in_file + 1, queue_state.get("next_id", 1))
+        gr.Info(f"Loaded {len(newly_loaded_queue)} tasks ({loaded_image_count} images).")
+        if error_messages: gr.Warning(" ".join(error_messages))
+    except Exception as e: print(f"Error loading queue: {e}"); traceback.print_exc(); gr.Warning(f"Failed to load queue: {str(e)[:200]}")
+    finally:
+        if uploaded_zip_file_obj and hasattr(uploaded_zip_file_obj, 'name') and uploaded_zip_file_obj.name and tempfile.gettempdir() in os.path.abspath(uploaded_zip_file_obj.name):
+            try: os.remove(uploaded_zip_file_obj.name)
+            except OSError: pass
+    return state_dict_gr_state, update_queue_df_display(queue_state)
+
+def autosave_queue_on_exit_action(state_dict_gr_state_ref):
+    """Saves the queue to a zip file on application exit."""
+    print("Attempting to autosave queue on exit...")
+    queue_state = get_queue_state(state_dict_gr_state_ref)
+    if not queue_state.get("queue"): print("Autosave: Queue is empty."); return
+    try:
+        _dummy_state_ignored, zip_b64_for_save = save_queue_to_zip(state_dict_gr_state_ref)
+        if zip_b64_for_save:
+            with open(AUTOSAVE_FILENAME, "wb") as f: f.write(base64.b64decode(zip_b64_for_save))
+            print(f"Autosave successful: Queue saved to {AUTOSAVE_FILENAME}")
+        else: print("Autosave failed: Could not generate zip data.")
+    except Exception as e: print(f"Error during autosave: {e}"); traceback.print_exc()
+
+def autoload_queue_on_start_action(state_dict_gr_state):
+    """Loads a previously autosaved queue when the application starts."""
+    queue_state = get_queue_state(state_dict_gr_state)
+    df_update = update_queue_df_display(queue_state)
+
+    if not queue_state["queue"] and Path(AUTOSAVE_FILENAME).is_file():
+        print(f"Autoloading queue from {AUTOSAVE_FILENAME}...")
+        class MockFilepath:
+            def __init__(self, name): self.name = name
+
+        temp_state_for_load = {"queue_state": queue_state.copy()}
+        loaded_state_result, df_update_after_load = load_queue_from_zip(temp_state_for_load, MockFilepath(AUTOSAVE_FILENAME))
+
+        if loaded_state_result["queue_state"]["queue"]:
+            queue_state.update(loaded_state_result["queue_state"])
+            df_update = df_update_after_load
+            print(f"Autoload successful. Loaded {len(queue_state['queue'])} tasks.")
+            try:
+                os.remove(AUTOSAVE_FILENAME)
+                print(f"Removed autosave file: {AUTOSAVE_FILENAME}")
+            except OSError as e:
+                print(f"Error removing autosave file '{AUTOSAVE_FILENAME}': {e}")
+        else:
+            print("Autoload: File existed but queue remains empty. Resetting queue.")
+            queue_state["queue"] = []
+            queue_state["next_id"] = 1
+            df_update = update_queue_df_display(queue_state)
+
+    is_processing_on_load = queue_state.get("processing", False) and bool(queue_state.get("queue"))
+    initial_video_path = state_dict_gr_state.get("last_completed_video_path")
+    if initial_video_path and not os.path.exists(initial_video_path):
+        print(f"Warning: Last completed video file not found at {initial_video_path}. Clearing reference.")
+        initial_video_path = None
+        state_dict_gr_state["last_completed_video_path"] = None
+
+    return (state_dict_gr_state, df_update, gr.update(interactive=not is_processing_on_load), gr.update(interactive=is_processing_on_load), gr.update(value=initial_video_path))
+
+def process_task_queue_main_loop(state_dict_gr_state):
+    """The main loop that processes tasks from the queue one by one."""
+    queue_state = get_queue_state(state_dict_gr_state)
+    shared_state.abort_event.clear()
+
+    output_stream_for_ui = state_dict_gr_state.get("active_output_stream_queue")
+
+    if queue_state["processing"]:
+        gr.Info("Queue processing is already active. Attempting to re-attach UI to live updates...")
+        if output_stream_for_ui is None:
+            gr.Warning("No active stream found in state. Queue processing may have been interrupted. Please clear queue or restart."); queue_state["processing"] = False
+            yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True)); return
+
+        # --- MODIFIED: Initial yield for re-attachment path ---
+        # Provide placeholder updates to progress/preview UI elements
+        yield (
+            state_dict_gr_state,
+            update_queue_df_display(queue_state),
+            gr.update(value=state_dict_gr_state.get("last_completed_video_path", None)), # Keep last video if available
+            gr.update(visible=True, value=None), # Make preview visible but start blank until new data
+            gr.update(value=f"Re-attaching to processing Task {queue_state['queue'][0]['id']}... Awaiting next preview."),
+            gr.update(value="<div style='text-align: center;'>Re-connecting...</div>"), # Generic HTML for progress bar
+            gr.update(interactive=False), # Process Queue button
+            gr.update(interactive=True),  # Abort button
+            gr.update(interactive=True)   # Reset button
+        )
+    elif not queue_state["queue"]:
+        gr.Info("Queue is empty. Add tasks to process.")
+        yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True)); return
+    else:
+        queue_state["processing"] = True
+        output_stream_for_ui = AsyncStream()
+        state_dict_gr_state["active_output_stream_queue"] = output_stream_for_ui
+        yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(visible=False), gr.update(value="Queue processing started..."), gr.update(value=""), gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True))
+
+    actual_output_queue = output_stream_for_ui.output_queue if output_stream_for_ui else None
+    if not actual_output_queue:
+        gr.Warning("Internal error: Output queue not available. Aborting."); queue_state["processing"] = False; state_dict_gr_state["active_output_stream_queue"] = None
+        yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True)); return
+
+    while queue_state["queue"] and not shared_state.abort_event.is_set():
+        with shared_state.queue_lock:
+            if not queue_state["queue"]: break
+            current_task_obj = queue_state["queue"][0]
+            task_parameters_for_worker = current_task_obj["params"]
+            current_task_id = current_task_obj["id"]
+
+        if task_parameters_for_worker.get('input_image') is None:
+            print(f"Skipping task {current_task_id}: Missing input image data.")
+            gr.Warning(f"Task {current_task_id} skipped: Input image is missing.")
+            with shared_state.queue_lock: 
+                current_task_obj["status"] = "error"; current_task_obj["error_message"] = "Missing Image"
+            yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(visible=False), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True)); break
+
+        if task_parameters_for_worker.get('seed') == -1: task_parameters_for_worker['seed'] = np.random.randint(0, 2**32 - 1)
+
+        print(f"Starting task {current_task_id} (Prompt: {task_parameters_for_worker.get('prompt', '')[:30]}...).")
+        current_task_obj["status"] = "processing"
+        yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(), gr.update(visible=False), gr.update(value=f"Processing Task {current_task_id}..."), gr.update(value=""), gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True))
+
+        worker_args = {
+            **task_parameters_for_worker,
+            'task_id': current_task_id, 'output_queue_ref': actual_output_queue, 'abort_event': shared_state.abort_event,
+            **shared_state.models
+        }
+        async_run(worker, **worker_args)
+
+        last_known_output_filename = state_dict_gr_state.get("last_completed_video_path", None)
+        task_completed_successfully = False
+        while True:
+            flag, data_from_worker = actual_output_queue.next()
+            if flag == 'progress':
+                msg_task_id, preview_np_array, desc_str, html_str = data_from_worker
+                if msg_task_id == current_task_id: yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value=last_known_output_filename), gr.update(visible=(preview_np_array is not None), value=preview_np_array), desc_str, html_str, gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True))
+            elif flag == 'file':
+                msg_task_id, segment_file_path, segment_info = data_from_worker
+                if msg_task_id == current_task_id: last_known_output_filename = segment_file_path; gr.Info(f"Task {current_task_id}: {segment_info}")
+                yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value=last_known_output_filename), gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True))
+            elif flag == 'aborted': current_task_obj["status"] = "aborted"; task_completed_successfully = False; break
+            elif flag == 'error': _, error_message_str = data_from_worker; gr.Warning(f"Task {current_task_id} Error: {error_message_str}"); current_task_obj["status"] = "error"; current_task_obj["error_message"] = str(error_message_str)[:100]; task_completed_successfully = False; break
+            elif flag == 'end': _, success_bool, final_video_path = data_from_worker; task_completed_successfully = success_bool; last_known_output_filename = final_video_path if success_bool else last_known_output_filename; current_task_obj["status"] = "done" if success_bool else "error"; break
+
+        with shared_state.queue_lock:
+            if queue_state["queue"] and queue_state["queue"][0]["id"] == current_task_id: queue_state["queue"].pop(0)
+        state_dict_gr_state["last_completed_video_path"] = last_known_output_filename if task_completed_successfully else None
+        final_desc = f"Task {current_task_id} {'completed' if task_completed_successfully else 'finished with issues'}."
+        yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value=state_dict_gr_state["last_completed_video_path"]), gr.update(visible=False), gr.update(value=final_desc), gr.update(value=""), gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=True))
+        if shared_state.abort_event.is_set(): gr.Info("Queue processing halted by user."); break
+
+    queue_state["processing"] = False; state_dict_gr_state["active_output_stream_queue"] = None
+    final_status_msg = "All tasks processed." if not shared_state.abort_event.is_set() else "Queue processing aborted."
+    yield (state_dict_gr_state, update_queue_df_display(queue_state), gr.update(value=state_dict_gr_state["last_completed_video_path"]), gr.update(visible=False), gr.update(value=final_status_msg), gr.update(value=""), gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True))
+
+def abort_current_task_processing_action(state_dict_gr_state):
+    """Sends the abort signal to the currently processing task."""
+    queue_state = get_queue_state(state_dict_gr_state)
+    if queue_state["processing"]:
+        gr.Info("Abort signal sent. Current task will attempt to stop shortly.")
+        shared_state.abort_event.set()
+    else:
+        gr.Info("Nothing is currently processing.")
+    return state_dict_gr_state, gr.update(interactive=not queue_state["processing"])

ui/shared_state.py

@@ -0,0 +1,60 @@
+# ui/shared_state.py
+import threading
+
+# --- Application-wide Threading Controls ---
+# Lock for thread-safe queue modifications.
+queue_lock = threading.Lock()
+# Event to signal the abortion of the current processing task.
+abort_event = threading.Event()
+
+# --- Model and Global State Containers ---
+# This dictionary will be populated at runtime by the main script after the models are loaded.
+# It allows other modules to access the models without circular imports.
+models = {}
+
+# This dictionary holds the application state, which is passed to the atexit
+# handler to enable the autosave functionality on browser close or exit.
+global_state_for_autosave = {}
+
+
+# --- UI and Parameter Mapping Constants ---
+# These constants define the structure of the UI and how UI components
+# map to the parameters of the backend generation worker.
+
+# Creative "Recipe" Parameters (for portable PNG metadata and task editing)
+CREATIVE_PARAM_KEYS = [
+    'prompt', 'n_prompt', 'total_second_length', 'seed', 'preview_frequency_ui',
+    'segments_to_decode_csv', 'gs_ui', 'gs_schedule_shape_ui', 'gs_final_ui', 'steps', 'cfg', 'rs'
+]
+
+# Environment/Debug Parameters (for the full workspace, machine/session-specific)
+ENVIRONMENT_PARAM_KEYS = [
+    'use_teacache', 'use_fp32_transformer_output_ui', 'gpu_memory_preservation',
+    'mp4_crf', 'output_folder_ui', 'latent_window_size'
+]
+
+# A comprehensive list of all UI components that define a task's parameters.
+ALL_TASK_UI_KEYS = CREATIVE_PARAM_KEYS + ENVIRONMENT_PARAM_KEYS
+
+# This maps the string keys of the Gradio UI components to the keyword argument
+# names expected by the 'worker' function in generation_core.py.
+UI_TO_WORKER_PARAM_MAP = {
+    'prompt': 'prompt',
+    'n_prompt': 'n_prompt',
+    'total_second_length': 'total_second_length',
+    'seed': 'seed',
+    'use_teacache': 'use_teacache',
+    'preview_frequency_ui': 'preview_frequency',
+    'segments_to_decode_csv': 'segments_to_decode_csv',
+    'gs_ui': 'gs',
+    'gs_schedule_shape_ui': 'gs_schedule_active',
+    'gs_final_ui': 'gs_final',
+    'steps': 'steps',
+    'cfg': 'cfg',
+    'latent_window_size': 'latent_window_size',
+    'gpu_memory_preservation': 'gpu_memory_preservation',
+    'use_fp32_transformer_output_ui': 'use_fp32_transformer_output',
+    'rs': 'rs',
+    'mp4_crf': 'mp4_crf',
+    'output_folder_ui': 'output_folder'
+}

ui/workspace.py

@@ -0,0 +1,206 @@
+# ui/workspace.py
+# Contains functions for saving, loading, and managing workspace settings from files.
+
+import gradio as gr
+import json
+import os
+import traceback
+import tkinter as tk
+from tkinter import filedialog
+from PIL import Image
+
+# Import shared state and other managers
+from . import shared_state
+from . import metadata as metadata_manager
+
+# --- Constants ---
+# Define filenames and the default output folder at the module level.
+# These were moved from demo_gradio_svc.py
+outputs_folder = './outputs_svc/'
+SETTINGS_FILENAME = "goan_settings.json"
+UNLOAD_SAVE_FILENAME = "goan_unload_save.json"
+REFRESH_IMAGE_FILENAME = "goan_refresh_image.png"
+
+# --- Core Save/Load Logic ---
+
+def get_default_values_map():
+    """Returns a dictionary with the default values for all UI settings."""
+    return {
+        'prompt': '', 'n_prompt': '', 'total_second_length': 5.0, 'seed': -1,
+        'use_teacache': True, 'preview_frequency_ui': 5, 'segments_to_decode_csv': '',
+        'gs_ui': 10.0, 'gs_schedule_shape_ui': 'Off', 'gs_final_ui': 10.0, 'steps': 25,
+        'cfg': 1.0, 'latent_window_size': 9, 'gpu_memory_preservation': 6.0,
+        'use_fp32_transformer_output_ui': False, 'rs': 0.0, 'mp4_crf': 18,
+        'output_folder_ui': outputs_folder,
+    }
+
+def save_settings_to_file(filepath, *ui_values_tuple):
+    """Saves a tuple of UI values to a specified JSON file."""
+    settings_to_save = dict(zip(shared_state.ALL_TASK_UI_KEYS, ui_values_tuple))
+    try:
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(settings_to_save, f, indent=4)
+        gr.Info(f"Workspace saved to {filepath}")
+    except Exception as e:
+        gr.Warning(f"Error saving workspace: {e}")
+        traceback.print_exc()
+
+def load_settings_from_file(filepath, return_updates=True):
+    """Loads settings from a JSON file and returns Gradio updates or raw values."""
+    default_values = get_default_values_map()
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            loaded_settings = json.load(f)
+        gr.Info(f"Loaded workspace from {filepath}")
+    except Exception as e:
+        gr.Warning(f"Could not load workspace from {filepath}: {e}")
+        loaded_settings = {}
+
+    final_settings = {**default_values, **loaded_settings}
+    output_values = [final_settings.get(key, default_values.get(key)) for key in shared_state.ALL_TASK_UI_KEYS]
+
+    # Type correction to prevent errors when loading from JSON
+    for i, key in enumerate(shared_state.ALL_TASK_UI_KEYS):
+        try:
+            if key in ['seed', 'latent_window_size', 'steps', 'mp4_crf', 'preview_frequency_ui']:
+                output_values[i] = int(output_values[i])
+            elif key in ['total_second_length', 'cfg', 'gs_ui', 'rs', 'gpu_memory_preservation', 'gs_final_ui']:
+                output_values[i] = float(output_values[i])
+            elif key in ['use_teacache', 'use_fp32_transformer_output_ui']:
+                output_values[i] = bool(output_values[i])
+        except (ValueError, TypeError):
+            output_values[i] = default_values.get(key)
+
+    return [gr.update(value=v) for v in output_values] if return_updates else output_values
+
+# --- NEW FUNCTION FOR INITIAL OUTPUT FOLDER LOADING ---
+def get_initial_output_folder_from_settings():
+    """
+    Attempts to load the 'output_folder_ui' value from UNLOAD_SAVE_FILENAME or SETTINGS_FILENAME.
+    If not found or an error occurs, returns the module's default outputs_folder.
+    Expands user home directory paths (e.g., '~') if present.
+    """
+    # Use the module-level outputs_folder as the fallback default
+    default_output_folder_path = outputs_folder
+
+    # Prioritize the temporary unload save file, then the default settings file
+    filename_to_check = None
+    if os.path.exists(UNLOAD_SAVE_FILENAME):
+        filename_to_check = UNLOAD_SAVE_FILENAME
+    elif os.path.exists(SETTINGS_FILENAME):
+        filename_to_check = SETTINGS_FILENAME
+
+    if filename_to_check:
+        try:
+            with open(filename_to_check, 'r', encoding='utf-8') as f:
+                settings = json.load(f)
+            # If 'output_folder_ui' exists in the loaded settings, use it
+            if 'output_folder_ui' in settings:
+                # Expand user path (e.g., '~/.goan/outputs/') to its absolute form
+                return os.path.expanduser(settings['output_folder_ui'])
+        except Exception as e:
+            # Log a warning if loading fails, but proceed with default
+            print(f"Warning: Could not load 'output_folder_ui' from {filename_to_check} for initial path setup: {e}")
+            traceback.print_exc() # Print full traceback for deeper debugging if needed
+
+    # If no file found, or loading failed, or key not present, return the module's default
+    return default_output_folder_path
+
+# --- UI Handler Functions ---
+
+def save_workspace(*ui_values_tuple):
+    """Opens a file dialog to save the full workspace settings."""
+    root = tk.Tk()
+    root.withdraw()
+    file_path = filedialog.asksaveasfilename(
+        defaultextension=".json",
+        initialfile="goan_workspace.json",
+        filetypes=[("JSON files", "*.json")]
+    )
+    root.destroy()
+    if file_path:
+        save_settings_to_file(file_path, *ui_values_tuple)
+    else:
+        gr.Warning("Save cancelled.")
+
+def save_as_default_workspace(*ui_values_tuple):
+    """Saves the current UI settings as the default startup configuration."""
+    gr.Info(f"Saving current settings as default to {SETTINGS_FILENAME}")
+    save_settings_to_file(SETTINGS_FILENAME, *ui_values_tuple)
+
+def save_ui_and_image_for_refresh(*args_from_ui_controls_tuple):
+    """Saves UI state and the current image to temporary files for session recovery."""
+    gallery_list = args_from_ui_controls_tuple[0]
+    all_ui_values_tuple = args_from_ui_controls_tuple[1:]
+    full_params_map = dict(zip(shared_state.ALL_TASK_UI_KEYS, all_ui_values_tuple))
+    settings_to_save = full_params_map.copy()
+
+    if gallery_list and isinstance(gallery_list[0], (tuple, Image.Image)):
+        pil_image = gallery_list[0][0] if isinstance(gallery_list[0], tuple) else gallery_list[0]
+        try:
+            creative_params = {k: full_params_map.get(k) for k in shared_state.CREATIVE_PARAM_KEYS}
+            pil_image = metadata_manager.write_image_metadata(pil_image, creative_params)
+
+            # Use the output folder path from the UI settings
+            output_folder_path = full_params_map.get('output_folder_ui', outputs_folder)
+            refresh_image_path = os.path.join(output_folder_path, REFRESH_IMAGE_FILENAME)
+
+            pil_image.save(refresh_image_path)
+            settings_to_save["refresh_image_path"] = refresh_image_path
+        except Exception as e:
+            gr.Warning(f"Could not save refresh image: {e}")
+
+    # Save all settings to the unload file
+    save_settings_to_file(UNLOAD_SAVE_FILENAME, *settings_to_save.values())
+
+def load_workspace():
+    """Opens a file dialog to load a workspace from a JSON file."""
+    root = tk.Tk()
+    root.withdraw()
+    file_path = filedialog.askopenfilename(filetypes=[("JSON files", "*.json")])
+    root.destroy()
+    return load_settings_from_file(file_path) if file_path else [gr.update()] * len(shared_state.ALL_TASK_UI_KEYS)
+
+def load_workspace_on_start():
+    """
+    Loads settings on app startup, prioritizing a temporary session file,
+    then a default file, and finally falling back to hardcoded defaults.
+    """
+    image_path_to_load = None
+    settings_file = None
+
+    if os.path.exists(UNLOAD_SAVE_FILENAME):
+        settings_file = UNLOAD_SAVE_FILENAME
+        try:
+            with open(settings_file, 'r') as f:
+                settings = json.load(f)
+            # Check if the saved image path exists
+            if "refresh_image_path" in settings and os.path.exists(settings["refresh_image_path"]):
+                image_path_to_load = settings["refresh_image_path"]
+        except Exception:
+            pass # Ignore errors reading the temp file
+    elif os.path.exists(SETTINGS_FILENAME):
+        settings_file = SETTINGS_FILENAME
+
+    if settings_file:
+        print(f"Loading workspace from {settings_file}")
+        ui_updates = load_settings_from_file(settings_file)
+        if image_path_to_load:
+            gr.Info("Restoring UI state and image from previous session.")
+        if settings_file == UNLOAD_SAVE_FILENAME:
+            os.remove(UNLOAD_SAVE_FILENAME) # Clean up temp file
+        return [image_path_to_load] + ui_updates
+
+    print("No workspace file found. Using default values.")
+    default_vals = get_default_values_map()
+    return [None] + [default_vals[key] for key in shared_state.ALL_TASK_UI_KEYS]
+
+def load_image_from_path(image_path):
+    """Loads an image from a given path and deletes the temporary file."""
+    if image_path and os.path.exists(image_path):
+        try:
+            # The gallery component expects a list of (image, name) tuples.
+            return gr.update(value=[(Image.open(image_path), "refresh_image")])
+        finally:
+            os.remove(image_path) # Clean up the temp image
+    return gr.update(value=None)