"""Viggle-Animate's fixed conditioning and model information."""

import torch

from mmgp import offload
from shared.utils import files_locator as fl


VIGGLE_ARCHITECTURE = "viggle_animate"
VIGGLE_REPO_ID = "DeepBeepMeep/MiniMax-H3"
VIGGLE_ASSET_FOLDER = "viggle_animate"
VIGGLE_PROMPT_FILE = "viggle_animate_fixed_prompt.safetensors"

VIGGLE_INFOS = """## Viggle-Animate: replace a character using one edited video frame

Provide the original **Control Video** and one **Edited Reference Frame** from that same video. Viggle follows the video's motion and camera while applying the character's edited appearance. No pose extraction or mask is needed.

### Prepare the edited frame
1. Extract a clear frame in which the character is visible and unobstructed. Any frame can be used; it does not have to be the first frame.
2. Open that frame in an image editor or image-editing model and replace the character. Preserve the pose of this selected frame, including body orientation, hands and props. Keep the camera framing, background and lighting consistent.
3. Save at the same dimensions without cropping or reframing, then upload the result as the Edited Reference Frame.

Example instruction for your image editor: *Replace the dancer with a humanoid robot. Match the dancer's posture, body orientation and hand positions. Preserve the props, background, camera framing, lighting and image dimensions.*

The edited image is an appearance reference, not an imposed first frame. If you edit a frame with raised arms, keep the arms raised in that image; the generated video still follows the changing poses of the original footage.

### Generation
Model includes 3-step LoRA distillation and requires Euler Scheduler.

Each generation window is capped at **124 frames** (about 5.2 seconds at 24 fps). Longer videos use Sliding Window with a fixed 124-frame window and 18-frame overlap by default. Each window follows the corresponding section of the Control Video and reuses the previous output for continuity.

Viggle uses a fixed, pre-encoded prompt. **The generation prompt is ignored; do not write anything there.** Express the desired replacement through the edited image.

Image alignment matters. Large pose or framing changes in the edited image can degrade the animation. Lip-sync, interacting characters and shots with cuts remain challenging.

### Control Audio

- **No Input Audio** (default): use the original Viggle workflow. H3 generates its own audio track; the fixed prompt requests silence, but does not guarantee a digitally silent track.
- **Use Custom Audio**: upload an audio file in **Custom Audio** to guide generation.
- **Reuse Control Video Audio**: use the Control Video's soundtrack to guide generation. The video must contain an audio track.

With either audio input, H3 holds the supplied audio fixed while generating the video. This is experimental for Viggle: its released workflow was trained without source audio, and improved lip-sync is not guaranteed. Use audio synchronized to the Control Video; unrelated speech can conflict with the original mouth motion. Audio is read from the beginning and trimmed to the requested video duration. If it is shorter, the remaining audio is generated by H3. When the input covers the full clip, WanGP keeps the original soundtrack in the output.

[Viggle-Animate model and examples](https://huggingface.co/Viggle/Viggle-Animate)
"""


def load_fixed_prompt(filename):
    fixed = torch.nn.Module()
    fixed.register_buffer("prompt_embeds", torch.empty(1, 362, 5120, dtype=torch.bfloat16, device="meta"))
    fixed.register_buffer("text_token_tags", torch.empty(362, dtype=torch.int64, device="meta"))
    offload.load_model_data(fixed, fl.locate_file(filename), writable_tensors=False, default_dtype=torch.bfloat16)
    return fixed