# Copyright 2023 Natural Synthetics Inc. All rights reserved.
# 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
#

from dataclasses import dataclass
from typing import Union

import numpy as np
import torch

from diffusers.utils import (
    BaseOutput,
)

@dataclass
class HotshotPipelineXLOutput(BaseOutput):
    videos: Union[torch.Tensor, np.ndarray]
hotshot_xl/__init__.py
hotshotco/Hotshot-XL-Automatic1111
# Copyright 2023 Natural Synthetics Inc. All rights reserved.
# 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
#

import torch
import torch.nn as nn
from diffusers.models.resnet import Upsample2D, Downsample2D, LoRACompatibleConv
from einops import rearrange


class Upsample3D(Upsample2D):
    def forward(self, hidden_states, output_size=None, scale: float = 1.0):
        f = hidden_states.shape[2]
        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
        hidden_states = super(Upsample3D, self).forward(hidden_states, output_size, scale)
        return rearrange(hidden_states, "(b f) c h w -> b c f h w", f=f)


class Downsample3D(Downsample2D):

    def forward(self, hidden_states, scale: float = 1.0):
        f = hidden_states.shape[2]
        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
        hidden_states = super(Downsample3D, self).forward(hidden_states, scale)
        return rearrange(hidden_states, "(b f) c h w -> b c f h w", f=f)


class Conv3d(LoRACompatibleConv):
    def forward(self, hidden_states, scale: float = 1.0):
        f = hidden_states.shape[2]
        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
        hidden_states = super().forward(hidden_states, scale)
        return rearrange(hidden_states, "(b f) c h w -> b c f h w", f=f)


class ResnetBlock3D(nn.Module):
    def __init__(
            self,
            *,
            in_channels,
            out_channels=None,
            conv_shortcut=False,
            dropout=0.0,
            temb_channels=512,
            groups=32,
            groups_out=None,
            pre_norm=True,
            eps=1e-6,
            non_linearity="silu",
            time_embedding_norm="default",
            output_scale_factor=1.0,
            use_in_shortcut=None,
            conv_shortcut_bias: bool = True,
    ):
        super().__init__()
        self.pre_norm = pre_norm
        self.pre_norm = True
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.use_conv_shortcut = conv_shortcut
        self.time_embedding_norm = time_embedding_norm
        self.output_scale_factor = output_scale_factor

        if groups_out is None:
            groups_out = groups

        self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
        self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)

        if temb_channels is not None:
            if self.time_embedding_norm == "default":
                time_emb_proj_out_channels = out_channels
            elif self.time_embedding_norm == "scale_shift":
                time_emb_proj_out_channels = out_channels * 2
            else:
                raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")

            self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
        else:
            self.time_emb_proj = None

        self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
        self.dropout = torch.nn.Dropout(dropout)
        self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)

        assert non_linearity == "silu"

        self.nonlinearity = nn.SiLU()

        self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut

        self.conv_shortcut = None
        if self.use_in_shortcut:
            self.conv_shortcut = Conv3d(
                in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias
            )

    def forward(self, input_tensor, temb):
        hidden_states = input_tensor

        hidden_states = self.norm1(hidden_states)
        hidden_states = self.nonlinearity(hidden_states)

        hidden_states = self.conv1(hidden_states)

        if temb is not None:
            temb = self.nonlinearity(temb)
            temb = self.time_emb_proj(temb)[:, :, None, None, None]

        if temb is not None and self.time_embedding_norm == "default":
            hidden_states = hidden_states + temb

        hidden_states = self.norm2(hidden_states)

        if temb is not None and self.time_embedding_norm == "scale_shift":
            scale, shift = torch.chunk(temb, 2, dim=1)
            hidden_states = hidden_states * (1 + scale) + shift

        hidden_states = self.nonlinearity(hidden_states)

        hidden_states = self.dropout(hidden_states)
        hidden_states = self.conv2(hidden_states)

        if self.conv_shortcut is not None:
            input_tensor = self.conv_shortcut(input_tensor)

        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor

        return output_tensor

hotshot_xl/models/resnet.py
hotshotco/Hotshot-XL-Automatic1111
import torch
import torch.nn as nn
from typing import Optional
from .transformer_temporal import TransformerTemporal

class Block(nn.Module):
    def __init__(self, temporal_attentions: list):
        super().__init__()
        self.temporal_attentions = nn.ModuleList(temporal_attentions)


class HotshotXLTemporalLayers(nn.Module):
    def __init__(self):
        super().__init__()

        self.down_blocks = nn.ModuleList()
        self.up_blocks = nn.ModuleList()

        for channel in [320, 640, 1280]:
            blks = []
            for _ in range(2):
                blks.append(TransformerTemporal(
                    num_attention_heads=8,
                    attention_head_dim=channel // 8,
                    in_channels=channel,
                    cross_attention_dim=None,
            ))
            self.down_blocks.append(Block(blks))

        for channel in [1280, 640, 320]:
            blks = []
            for _ in range(3):
                blks.append(TransformerTemporal(
                    num_attention_heads=8,
                    attention_head_dim=channel // 8,
                    in_channels=channel,
                    cross_attention_dim=None,
            ))
            self.up_blocks.append(Block(blks))

    def get_temporal_layer(self, block_direction: int, block_index: int, attention_index: int):
        return (self.down_blocks if block_direction == -1 else self.up_blocks)[block_index].temporal_attentions[attention_index]


if __name__ == "__main__":
    layers = HotshotXLTemporalLayers()
    from safetensors import safe_open

    torch_model = {}
    with safe_open("/path/to/hsxl_temporal_layers.safetensors", framework="pt",
                   device="cuda") as f:
        for key in f.keys():
            torch_model[key] = f.get_tensor(key)

    layers.load_state_dict(torch_model)


hotshot_xl/models/temporal_layers.py
hotshotco/Hotshot-XL-Automatic1111
# Copyright 2023 Natural Synthetics Inc. All rights reserved.
# 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
#

from dataclasses import dataclass
from typing import Optional
import torch
from torch import nn
from diffusers.utils import BaseOutput
from diffusers.models.transformer_2d import Transformer2DModel
from einops import rearrange, repeat
from typing import Dict, Any

@dataclass
class Transformer3DModelOutput(BaseOutput):
    """
        The output of [`Transformer3DModel`].

        Args:
            sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`:
                The hidden states output conditioned on the `encoder_hidden_states` input.
        """

    sample: torch.FloatTensor


class Transformer3DModel(Transformer2DModel):

    def __init__(self, *args, **kwargs):
        super(Transformer3DModel, self).__init__(*args, **kwargs)
        nn.init.zeros_(self.proj_out.weight.data)
        nn.init.zeros_(self.proj_out.bias.data)

    def forward(
            self,
            hidden_states: torch.Tensor,
            encoder_hidden_states: Optional[torch.Tensor] = None,
            timestep: Optional[torch.LongTensor] = None,
            class_labels: Optional[torch.LongTensor] = None,
            cross_attention_kwargs: Dict[str, Any] = None,
            attention_mask: Optional[torch.Tensor] = None,
            encoder_attention_mask: Optional[torch.Tensor] = None,
            enable_temporal_layers: bool = True,
            positional_embedding: Optional[torch.Tensor] = None,
            return_dict: bool = True,
    ):

        is_video = len(hidden_states.shape) == 5

        if is_video:
            f = hidden_states.shape[2]
            hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
            encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=f)

        hidden_states = super(Transformer3DModel, self).forward(hidden_states,
                                                                encoder_hidden_states,
                                                                timestep,
                                                                class_labels,
                                                                cross_attention_kwargs,
                                                                attention_mask,
                                                                encoder_attention_mask,
                                                                return_dict=False)[0]

        if is_video:
            hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", f=f)

        if not return_dict:
            return (hidden_states,)

        return Transformer3DModelOutput(sample=hidden_states)
hotshot_xl/models/transformer_3d.py
hotshotco/Hotshot-XL-Automatic1111
# Copyright 2023 Natural Synthetics Inc. All rights reserved.
# 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
#

import torch
import math
from dataclasses import dataclass
from torch import nn
from diffusers.utils import BaseOutput
from diffusers.models.attention import Attention, FeedForward
from einops import rearrange, repeat
from typing import Optional


class PositionalEncoding(nn.Module):
    """
    Implements positional encoding as described in "Attention Is All You Need".
    Adds sinusoidal based positional encodings to the input tensor.
    """

    _SCALE_FACTOR = 10000.0  # Scale factor used in the positional encoding computation.

    def __init__(self, dim: int, dropout: float = 0.0, max_length: int = 24):
        super(PositionalEncoding, self).__init__()

        self.dropout = nn.Dropout(p=dropout)

        # The size is (1, max_length, dim) to allow easy addition to input tensors.
        positional_encoding = torch.zeros(1, max_length, dim)

        # Position and dim are used in the sinusoidal computation.
        position = torch.arange(max_length).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, dim, 2) * (-math.log(self._SCALE_FACTOR) / dim))

        positional_encoding[0, :, 0::2] = torch.sin(position * div_term)
        positional_encoding[0, :, 1::2] = torch.cos(position * div_term)

        # Register the positional encoding matrix as a buffer,
        # so it's part of the model's state but not the parameters.
        self.register_buffer('positional_encoding', positional_encoding)

    def forward(self, hidden_states: torch.Tensor, length: int) -> torch.Tensor:
        hidden_states = hidden_states + self.positional_encoding[:, :length]
        return self.dropout(hidden_states)


class TemporalAttention(Attention):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.pos_encoder = PositionalEncoding(kwargs["query_dim"], dropout=0)

    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, number_of_frames=8):
        sequence_length = hidden_states.shape[1]
        hidden_states = rearrange(hidden_states, "(b f) s c -> (b s) f c", f=number_of_frames)
        hidden_states = self.pos_encoder(hidden_states, length=number_of_frames)

        if encoder_hidden_states:
            encoder_hidden_states = repeat(encoder_hidden_states, "b n c -> (b s) n c", s=sequence_length)

        hidden_states = super().forward(hidden_states, encoder_hidden_states, attention_mask=attention_mask)

        return rearrange(hidden_states, "(b s) f c -> (b f) s c", s=sequence_length)


@dataclass
class TransformerTemporalOutput(BaseOutput):
    sample: torch.FloatTensor


class TransformerTemporal(nn.Module):
    def __init__(
            self,
            num_attention_heads: int,
            attention_head_dim: int,
            in_channels: int,
            num_layers: int = 1,
            dropout: float = 0.0,
            norm_num_groups: int = 32,
            cross_attention_dim: Optional[int] = None,
            attention_bias: bool = False,
            activation_fn: str = "geglu",
            upcast_attention: bool = False,
    ):
        super().__init__()

        inner_dim = num_attention_heads * attention_head_dim

        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
        self.proj_in = nn.Linear(in_channels, inner_dim)

        self.transformer_blocks = nn.ModuleList(
            [
                TransformerBlock(
                    dim=inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    dropout=dropout,
                    activation_fn=activation_fn,
                    attention_bias=attention_bias,
                    upcast_attention=upcast_attention,
                    cross_attention_dim=cross_attention_dim
                )
                for _ in range(num_layers)
            ]
        )
        self.proj_out = nn.Linear(inner_dim, in_channels)

    def forward(self, hidden_states, encoder_hidden_states=None):

        _, num_channels, f, height, width = hidden_states.shape
        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")

        skip = hidden_states

        hidden_states = self.norm(hidden_states)
        hidden_states = rearrange(hidden_states, "bf c h w -> bf (h w) c")
        hidden_states = self.proj_in(hidden_states)

        for block in self.transformer_blocks:
            hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states, number_of_frames=f)

        hidden_states = self.proj_out(hidden_states)
        hidden_states = rearrange(hidden_states, "bf (h w) c -> bf c h w", h=height, w=width).contiguous()

        output = hidden_states + skip
        output = rearrange(output, "(b f) c h w -> b c f h w", f=f)

        return output


class TransformerBlock(nn.Module):
    def __init__(
            self,
            dim,
            num_attention_heads,
            attention_head_dim,
            dropout=0.0,
            activation_fn="geglu",
            attention_bias=False,
            upcast_attention=False,
            depth=2,
            cross_attention_dim: Optional[int] = None
    ):
        super().__init__()

        self.is_cross = cross_attention_dim is not None

        attention_blocks = []
        norms = []

        for _ in range(depth):
            attention_blocks.append(
                TemporalAttention(
                    query_dim=dim,
                    cross_attention_dim=cross_attention_dim,
                    heads=num_attention_heads,
                    dim_head=attention_head_dim,
                    dropout=dropout,
                    bias=attention_bias,
                    upcast_attention=upcast_attention,
                )
            )
            norms.append(nn.LayerNorm(dim))

        self.attention_blocks = nn.ModuleList(attention_blocks)
        self.norms = nn.ModuleList(norms)

        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
        self.ff_norm = nn.LayerNorm(dim)

    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, number_of_frames=None):

        if not self.is_cross:
            encoder_hidden_states = None

        for block, norm in zip(self.attention_blocks, self.norms):
            norm_hidden_states = norm(hidden_states)
            hidden_states = block(
                norm_hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                attention_mask=attention_mask,
                number_of_frames=number_of_frames
            ) + hidden_states

        norm_hidden_states = self.ff_norm(hidden_states)
        hidden_states = self.ff(norm_hidden_states) + hidden_states

        output = hidden_states
        return output

hotshot_xl/models/transformer_temporal.py
hotshotco/Hotshot-XL-Automatic1111
# Copyright 2023 Natural Synthetics Inc. All rights reserved.
#
# 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.

from typing import List, Union
from io import BytesIO
import PIL
from PIL import ImageSequence, Image
import requests
import os
import numpy as np
import imageio
import hashlib


def get_image(img_path) -> PIL.Image.Image:
    if img_path.startswith("http"):
        return PIL.Image.open(requests.get(img_path, stream=True).raw)
    if os.path.exists(img_path):
        return Image.open(img_path)
    raise Exception("File not found")

def images_to_gif_bytes(images: List, duration: int = 1000) -> bytes:
    with BytesIO() as output_buffer:
        # Save the first image
        images[0].save(output_buffer,
                       format='GIF',
                       save_all=True,
                       append_images=images[1:],
                       duration=duration,
                       loop=0)  # 0 means the GIF will loop indefinitely

        # Get the byte array from the buffer
        gif_bytes = output_buffer.getvalue()

    return gif_bytes

def save_as_gif(images: List, file_path: str, duration: int = 1000):
    with open(file_path, "wb") as f:
        f.write(images_to_gif_bytes(images, duration))

def images_to_mp4_bytes(images: List[Image.Image], duration: int = 1000) -> bytes:
        with BytesIO() as output_buffer:
            with imageio.get_writer(output_buffer, format='mp4', fps=1/(duration/1000)) as writer:
                for img in images:
                    writer.append_data(np.array(img))
            mp4_bytes = output_buffer.getvalue()

        return mp4_bytes

def save_as_mp4(images: List[Image.Image], file_path: str, duration: int = 1000):
    with open(file_path, "wb") as f:
        f.write(images_to_mp4_bytes(images, duration))

def scale_aspect_fill(img, new_width, new_height):
    new_width = int(new_width)
    new_height = int(new_height)

    original_width, original_height = img.size
    ratio_w = float(new_width) / original_width
    ratio_h = float(new_height) / original_height

    if ratio_w > ratio_h:
        # It must be fixed by width
        resize_width = new_width
        resize_height = round(original_height * ratio_w)
    else:
        # Fixed by height
        resize_width = round(original_width * ratio_h)
        resize_height = new_height

    img_resized = img.resize((resize_width, resize_height), Image.LANCZOS)

    # Calculate cropping boundaries and do crop
    left = (resize_width - new_width) / 2
    top = (resize_height - new_height) / 2
    right = (resize_width + new_width) / 2
    bottom = (resize_height + new_height) / 2

    img_cropped = img_resized.crop((left, top, right, bottom))

    return img_cropped

def extract_gif_frames_from_midpoint(image: Union[str, PIL.Image.Image], fps: int=8, target_duration: int=1000) -> list:
    # Load the GIF
    image = get_image(image) if type(image) is str else image

    frames = []

    estimated_frame_time = None

    # some gifs contain the duration - others don't
    # so if there is a duration we will grab it otherwise we will fall back

    for frame in ImageSequence.Iterator(image):

        frames.append(frame.copy())
        if 'duration' in frame.info:
            frame_info_duration = frame.info['duration']
            if frame_info_duration > 0:
                estimated_frame_time = frame_info_duration

    if estimated_frame_time is None:
        if len(frames) <= 16:
            # assume it's 8fps
            estimated_frame_time = 1000 // 8
        else:
            # assume it's 15 fps
            estimated_frame_time = 70

    if len(frames) < fps:
        raise ValueError(f"fps of {fps} is too small for this gif as it only has {len(frames)} frames.")

    skip = len(frames) // fps
    upper_bound_index = len(frames) - 1

    best_indices = [x for x in range(0, len(frames), skip)][:fps]
    offset = int(upper_bound_index - best_indices[-1]) // 2
    best_indices = [x + offset for x in best_indices]
    best_duration = (best_indices[-1] - best_indices[0]) * estimated_frame_time

    while True:

        skip -= 1

        if skip == 0:
            break

        indices = [x for x in range(0, len(frames), skip)][:fps]

        # center the indices, so we sample the middle of the gif...
        offset = int(upper_bound_index - indices[-1]) // 2
        if offset == 0:
            # can't shift
            break
        indices = [x + offset for x in indices]

        # is the new duration closer to the target than last guess?
        duration = (indices[-1] - indices[0]) * estimated_frame_time
        if abs(duration - target_duration) > abs(best_duration - target_duration):
            break

        best_indices = indices
        best_duration = duration

    return [frames[index] for index in best_indices]

def get_crop_coordinates(old_size: tuple, new_size: tuple) -> tuple:
    """
    Calculate the crop coordinates after scaling an image to fit a new size.

    :param old_size: tuple of the form (width, height) representing the original size of the image.
    :param new_size: tuple of the form (width, height) representing the desired size after scaling.
    :return: tuple of the form (left, upper, right, lower) representing the normalized crop coordinates.
    """
    # Check if the input tuples have the right form (width, height)
    if not (isinstance(old_size, tuple) and isinstance(new_size, tuple) and
            len(old_size) == 2 and len(new_size) == 2):
        raise ValueError("old_size and new_size should be tuples of the form (width, height)")

    # Extract the width and height from the old and new sizes
    old_width, old_height = old_size
    new_width, new_height = new_size

    # Calculate the ratios for width and height
    ratio_w = float(new_width) / old_width
    ratio_h = float(new_height) / old_height

    # Determine which dimension is fixed (width or height)
    if ratio_w > ratio_h:
        # It must be fixed by width
        resize_width = new_width
        resize_height = round(old_height * ratio_w)
    else:
        # Fixed by height
        resize_width = round(old_width * ratio_h)
        resize_height = new_height

    # Calculate cropping boundaries in the resized image space
    left = (resize_width - new_width) / 2
    upper = (resize_height - new_height) / 2
    right = (resize_width + new_width) / 2
    lower = (resize_height + new_height) / 2

    # Normalize the cropping coordinates

    # Return the normalized coordinates as a tuple
    return (left, upper, right, lower)

aspect_ratio_to_1024_map = {
    "0.42": [640,  1536],
    "0.57": [768,  1344],
    "0.68": [832,  1216],
    "1.00": [1024, 1024],
    "1.46": [1216,  832],
    "1.75": [1344,  768],
    "2.40": [1536,  640]
}

res_to_aspect_map = {
    1024: aspect_ratio_to_1024_map,
    512: {key: [value[0] // 2, value[1] // 2] for key, value in aspect_ratio_to_1024_map.items()},
}

def best_aspect_ratio(aspect_ratio: float, resolution: int):

    map = res_to_aspect_map[resolution]

    d = 99999999
    res = None
    for key, value in map.items():
        ar = value[0] / value[1]
        diff = abs(aspect_ratio - ar)
        if diff < d:
            d = diff
            res = value

    ar = res[0] / res[1]
    return f"{ar:.2f}", res

def hash_str(input: str):
    hash_object = hashlib.sha1(input.encode("utf-8"))
    return hash_object.hexdigest()
hotshot_xl/utils.py
hotshotco/Hotshot-XL-Automatic1111
print("Hotshot-XL install.py")

import logging
import os
from modules import scripts
import packages

os.makedirs(os.path.join(scripts.basedir(), "model"), exist_ok=True)

logger = logging.getLogger()

# if not packages.is_diffusers_installed():
#     packages.install_diffusers()
install.py
hotshotco/Hotshot-XL-Automatic1111
import subprocess

min_diffusers_version = "0.21.4"

# Depending on the Python version, we might need the external package or the standard library one
try:
    from importlib.metadata import version, PackageNotFoundError
except ImportError:
    from importlib_metadata import version, PackageNotFoundError

def is_diffusers_installed():
    try:
        # Check if 'diffusers' is installed
        installed_version = version("diffusers")
        return installed_version
    except PackageNotFoundError:
        return None

def install_diffusers():
    installed_version = is_diffusers_installed()

    if installed_version is None:
        print("Diffusers package is not installed. Installing now...")
        subprocess.run(f"pip install diffusers=={min_diffusers_version} --upgrade".split(" "))
    elif tuple(map(int, installed_version.split('.'))) < tuple(map(int, min_diffusers_version.split('.'))):
        print(f"Diffusers package version {installed_version} is less than required {min_diffusers_version}. Upgrading now...")
        subprocess.run(f"pip install diffusers=={min_diffusers_version} --upgrade".split(" "))
    else:
        print(f"Diffusers package version {installed_version} is up to date.")


packages.py
hotshotco/Hotshot-XL-Automatic1111
import install

import os
import gradio as gr
from modules import script_callbacks, scripts, shared
from modules.processing import (Processed, StableDiffusionProcessing,
                                StableDiffusionProcessingImg2Img)
from typing import Any, Union, Dict

try:
    from scripts.hotshot_xl_ui import HotshotXLUiGroup, HotshotXLParams
    from scripts.hotshot_xl_model_controller import model_controller
    from scripts.hotshot_xl_output import HotshotXLOutput
except:
    ...

script_dir = scripts.basedir()

class HotshotXLScript(scripts.Script):

    def __init__(self):
        print("HotshotXLScript init")
        self.lora_hacker = None
        self.cfg_hacker = None
        self.cn_hacker = None
        os.makedirs(self.model_directory, exist_ok=True)

    @property
    def model_directory(self):
        return shared.opts.data.get("hotshot_xl_model_path", os.path.join(script_dir, "model"))

    def title(self):
        return "Hotshot-XL"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        return (HotshotXLUiGroup().render(is_img2img, self.model_directory),)

    def before_process(
            self, p: StableDiffusionProcessing, params: Union[Dict, HotshotXLParams]
    ):
        if shared.sd_model and not shared.sd_model.is_sdxl:
            print("disabling because sdxl is not loaded...")
            return

        if isinstance(params, dict): params = HotshotXLParams(**params)
        if params.enable:
            params.set_p(p)

            # we don't cache the conditioning because we modify it using
            # original width and height & target width and height
            # which aren't part of the cache key

            p.cached_c = [None, None]
            p.cached_uc = [None, None]

            model_path = os.path.join(self.model_directory, params.model)
            model_controller.load_and_inject(
                shared.sd_model,
                model_path,
                cond={
                    "positive": {
                        "og_size": (params.original_size_width, params.original_size_height),
                        "tgt_size": (params.target_size_width, params.target_size_width),
                    },
                    "negative": {
                        "og_size": (params.negative_original_size_width, params.negative_original_size_height),
                        "tgt_size": (params.negative_target_size_width, params.negative_target_size_width),
                    }
                }
            )
            model_controller.set_video_length(params.video_length)

    def before_process_batch(
            self, p: StableDiffusionProcessing, params: Union[Dict, HotshotXLParams], **kwargs
    ):
        if isinstance(params, dict): params = HotshotXLParams(**params)
        if params.enable and isinstance(p, StableDiffusionProcessingImg2Img):
            # - not supported
            ...

    def postprocess(
            self, p: StableDiffusionProcessing, res: Processed, params: Union[Dict, HotshotXLParams]
    ):
        if isinstance(params, dict): params = HotshotXLParams(**params)
        if params.enable:
            model_controller.restore(shared.sd_model)
            HotshotXLOutput().output(p, res, params)


def on_ui_settings():
    section = ("hotshotxl", "Hotshot-XL")
    pass

def on_model_load(model):
    print("model loaded. ")

script_callbacks.on_ui_settings(on_ui_settings)
script_callbacks.on_after_component(HotshotXLUiGroup.on_after_component)
script_callbacks.on_model_loaded(on_model_load)


scripts/hotshot_xl.py
hotshotco/Hotshot-XL-Automatic1111
import gc
from dataclasses import dataclass
from typing import Optional

import sgm.models.diffusion
import torch
import torch.nn as nn
from einops import rearrange
from hotshot_xl.models.temporal_layers import HotshotXLTemporalLayers
from hotshot_xl.utils import hash_str
from safetensors import safe_open
from sgm.modules.attention import SpatialTransformer
from sgm.modules.diffusionmodules.openaimodel import ResBlock, TimestepEmbedSequential
from sgm.modules.diffusionmodules.util import GroupNorm32

from modules import sd_models_xl
from modules import devices, shared, prompt_parser
from modules.devices import torch_gc
from modules.processing import StableDiffusionProcessingTxt2Img

@dataclass
class TemporalModel:
    model: HotshotXLTemporalLayers
    model_hash: str

class TimeCentricTensorReshaper(nn.Module):
    def __init__(self, module):
        super().__init__()
        self.module = module
        self.video_length = 8

    def forward(self, hidden_states, encoder_hidden_states=None):

        if self.video_length < 1:
            # skip this module
            return hidden_states

        hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", f=self.video_length)
        hidden_states = self.module(hidden_states, encoder_hidden_states)
        return rearrange(hidden_states, "b c f h w -> (b f) c h w")


class HotshotXLModelController:

    def __init__(self):
        self.current_loaded_temporal_layers: Optional[TemporalModel] = None
        self.time_centric_tensor_reshaper_cache = []

    def set_video_length(self, video_length: int = 1):
        if video_length < 1:
            print("Warning - video length is < 1 - temporal layers will be deactivated.")

        for module in self.time_centric_tensor_reshaper_cache:
            module.video_length = video_length

    def load_and_inject(self,
                        sd_model,
                        temporal_layers_model_path: str,
                        cond: dict
                        ):

        if (not self.current_loaded_temporal_layers or
                self.current_loaded_temporal_layers.model_hash != hash_str(temporal_layers_model_path)):
            temporal_layers = HotshotXLTemporalLayers()
            torch_model = {}
            with safe_open(temporal_layers_model_path, framework="pt", device="cuda") as f:
                for key in f.keys():
                    torch_model[key] = f.get_tensor(key)
            temporal_layers.load_state_dict(torch_model)
            self.current_loaded_temporal_layers = TemporalModel(
                model=temporal_layers.to(device=sd_model.device, dtype=sd_model.dtype),
                model_hash=temporal_layers_model_path
            )

        self._hijack_sdxl_model(sd_model, self.current_loaded_temporal_layers.model)
        self._hijack_sd_models_xl_conditioning(cond)

    def _inject(self, spatial_module, temporal_module, type_to_insert_after: type):
        for i, module in enumerate(spatial_module):
            if type(module) == type_to_insert_after:
                self.time_centric_tensor_reshaper_cache.append(
                    TimeCentricTensorReshaper(temporal_module)
                )
                spatial_module.insert(i + 1, self.time_centric_tensor_reshaper_cache[-1])
                break

    def _hijack_sd_models_xl_conditioning(self, cond: dict):
        self.old_get_learned_conditioning = sd_models_xl.get_learned_conditioning

        # auto 1111 is limited in controlling the conditioning for sdxl
        # hotshot xl infernce takes advantage of the original size and target size
        # to give the best performance

        def get_learned_conditioning(self: sgm.models.diffusion.DiffusionEngine,
                                     batch: prompt_parser.SdConditioning | list[str]):
            for embedder in self.conditioner.embedders:
                embedder.ucg_rate = 0.0

            is_negative_prompt = getattr(batch, 'is_negative_prompt', False)

            c = cond['negative'] if is_negative_prompt else cond['positive']

            aesthetic_score = shared.opts.sdxl_refiner_low_aesthetic_score if is_negative_prompt else shared.opts.sdxl_refiner_high_aesthetic_score

            devices_args = dict(device=devices.device, dtype=devices.dtype)

            sdxl_conds = {
                "txt": batch,
                "original_size_as_tuple": torch.tensor(list(c['og_size']), **devices_args).repeat(len(batch), 1),
                "crop_coords_top_left": torch.tensor([shared.opts.sdxl_crop_top, shared.opts.sdxl_crop_left],
                                                     **devices_args).repeat(len(batch), 1),
                "target_size_as_tuple": torch.tensor(list(c['tgt_size']), **devices_args).repeat(len(batch), 1),
                "aesthetic_score": torch.tensor([aesthetic_score], **devices_args).repeat(len(batch), 1),
            }

            force_zero_negative_prompt = is_negative_prompt and all(x == '' for x in batch)
            c = self.conditioner(sdxl_conds, force_zero_embeddings=['txt'] if force_zero_negative_prompt else [])

            return c

        sgm.models.diffusion.DiffusionEngine.get_learned_conditioning = get_learned_conditioning

    def _hijack_sdxl_model(self, sd_model, temporal_layers: HotshotXLTemporalLayers):
        unet = sd_model.model.diffusion_model
        setattr(unet, "hsxl_hijacked", True)

        unet_input_block_index_to_temporal_layer_map = {
            1: {"block_index": 0, "attention_index": 0, "insert_after": ResBlock},
            2: {"block_index": 0, "attention_index": 1, "insert_after": ResBlock},

            4: {"block_index": 1, "attention_index": 0, "insert_after": SpatialTransformer},
            5: {"block_index": 1, "attention_index": 1, "insert_after": SpatialTransformer},

            7: {"block_index": 2, "attention_index": 0, "insert_after": SpatialTransformer},
            8: {"block_index": 2, "attention_index": 1, "insert_after": SpatialTransformer}
        }

        unet_output_block_index_to_temporal_layer_map = {
            0: {"block_index": 0, "attention_index": 0, "insert_after": SpatialTransformer},
            1: {"block_index": 0, "attention_index": 1, "insert_after": SpatialTransformer},
            2: {"block_index": 0, "attention_index": 2, "insert_after": SpatialTransformer},

            3: {"block_index": 1, "attention_index": 0, "insert_after": SpatialTransformer},
            4: {"block_index": 1, "attention_index": 1, "insert_after": SpatialTransformer},
            5: {"block_index": 1, "attention_index": 2, "insert_after": SpatialTransformer},

            6: {"block_index": 2, "attention_index": 0, "insert_after": ResBlock},
            7: {"block_index": 2, "attention_index": 1, "insert_after": ResBlock},
            8: {"block_index": 2, "attention_index": 2, "insert_after": ResBlock},
        }

        for unet_block_index, temporal_info in unet_input_block_index_to_temporal_layer_map.items():
            temporal_module = temporal_layers.get_temporal_layer(
                -1,
                temporal_info['block_index'],
                temporal_info['attention_index']
            )

            self._inject(unet.input_blocks[unet_block_index], temporal_module, temporal_info['insert_after'])

        for unet_block_index, temporal_info in unet_output_block_index_to_temporal_layer_map.items():
            temporal_module = temporal_layers.get_temporal_layer(
                1,
                temporal_info['block_index'],
                temporal_info['attention_index']
            )

            self._inject(unet.output_blocks[unet_block_index], temporal_module, temporal_info['insert_after'])

        # hotshot xl group norms with a different tensor arrangement - same as animate diff

        self.gn32_original_forward = GroupNorm32.forward
        gn32_original_forward = self.gn32_original_forward

        def groupnorm32_mm_forward(self, x):
            x = rearrange(x, "(b f) c h w -> b c f h w", b=2)
            x = gn32_original_forward(self, x)
            x = rearrange(x, "b c f h w -> (b f) c h w", b=2)
            return x

        GroupNorm32.forward = groupnorm32_mm_forward

    def unload(self):

        if (shared.sd_model.model.diffusion_model
                and hasattr(shared.sd_model.model.diffusion_model, "hsxl_hijacked")
                and shared.sd_model.model.diffusion_model.hsxl_hijacked is True):

            self.restore(shared.sd_model)

        self.current_loaded_temporal_layers = None
        torch_gc()
        gc.collect()

    def restore(self, sd_model):

        sgm.models.diffusion.DiffusionEngine.get_learned_conditioning = self.old_get_learned_conditioning
        unet = sd_model.model.diffusion_model
        GroupNorm32.forward = self.gn32_original_forward

        self.time_centric_tensor_reshaper_cache = []

        for block in unet.input_blocks:
            if type(block) is TimestepEmbedSequential:
                for i, module in enumerate(block):
                    if type(module) is TimeCentricTensorReshaper:
                        block.pop(i)
                        break

        for block in unet.output_blocks:
            if type(block) is TimestepEmbedSequential:
                for i, module in enumerate(block):
                    if type(module) is TimeCentricTensorReshaper:
                        block.pop(i)
                        break

        setattr(unet, "hsxl_hijacked", False)

model_controller = HotshotXLModelController()
scripts/hotshot_xl_model_controller.py
hotshotco/Hotshot-XL-Automatic1111
import base64
from pathlib import Path

import imageio.v3 as imageio
import numpy as np
from PIL import Image
from modules import images, shared
from modules.processing import Processed, StableDiffusionProcessing
from io import BytesIO
import logging
logger = logger = logging.getLogger()
from scripts.hotshot_xl_ui import HotshotXLParams


class HotshotXLOutput:

    working_dir = "HotshotXl"
    def output(
            self, p: StableDiffusionProcessing, res: Processed, params: HotshotXLParams
    ):
        video_paths = []
        logger.info("Merging images into GIF.")
        Path(f"{p.outpath_samples}/{self.working_dir}").mkdir(exist_ok=True, parents=True)
        step = params.video_length if params.video_length > params.batch_size else params.batch_size
        for i in range(res.index_of_first_image, len(res.images), step):
            # frame interpolation replaces video_list with interpolated frames
            # so make a copy instead of a slice (reference), to avoid modifying res
            video_list = [image.copy() for image in res.images[i: i + params.video_length]]

            seq = images.get_next_sequence_number(
                f"{p.outpath_samples}/{self.working_dir}", ""
            )
            filename = f"{seq:05}-{res.seed}"
            video_path_prefix = f"{p.outpath_samples}/{self.working_dir}/{filename}."

            video_list = self._add_reverse(params, video_list)
            #video_list = self._interp(p, params, video_list, filename)
            video_paths += self._save(params, video_list, video_path_prefix, res, i)

        if len(video_paths) > 0:
            if not p.is_api:
                res.images = video_paths
            # else:
            #     res.images = self._encode_video_to_b64(video_paths)

    def _add_reverse(self, params: HotshotXLParams, video_list: list):
        if 0 in params.reverse:
            video_list_reverse = video_list[::-1]
            if 1 in params.reverse:
                video_list_reverse.pop(0)
            if 2 in params.reverse:
                video_list_reverse.pop(-1)
            return video_list + video_list_reverse
        return video_list

    def _interp(
            self,
            p: StableDiffusionProcessing,
            params: HotshotXLParams,
            video_list: list,
            filename: str
    ):
        if params.interp not in ['FILM']:
            return video_list

        try:
            from deforum_helpers.frame_interpolation import (
                calculate_frames_to_add, check_and_download_film_model)
            from film_interpolation.film_inference import run_film_interp_infer
        except ImportError:
            logger.error(
                "Deforum not found. Please install: https://github.com/deforum-art/deforum-for-automatic1111-webui.git")
            return video_list

        import glob
        import os
        import shutil

        import modules.paths as ph

        # load film model
        deforum_models_path = ph.models_path + '/Deforum'
        film_model_folder = os.path.join(deforum_models_path, 'film_interpolation')
        film_model_name = 'film_net_fp16.pt'
        film_model_path = os.path.join(film_model_folder, film_model_name)
        check_and_download_film_model('film_net_fp16.pt', film_model_folder)

        film_in_between_frames_count = calculate_frames_to_add(len(video_list), params.interp_x)

        # save original frames to tmp folder for deforum input
        tmp_folder = f"{p.outpath_samples}/{self.working_dir}/tmp"
        input_folder = f"{tmp_folder}/input"
        os.makedirs(input_folder, exist_ok=True)
        for tmp_seq, frame in enumerate(video_list):
            imageio.imwrite(f"{input_folder}/{tmp_seq:05}.png", frame)

        # deforum saves output frames to tmp/{filename}
        save_folder = f"{tmp_folder}/{filename}"
        os.makedirs(save_folder, exist_ok=True)

        run_film_interp_infer(
            model_path=film_model_path,
            input_folder=input_folder,
            save_folder=save_folder,
            inter_frames=film_in_between_frames_count)

        # load deforum output frames and replace video_list
        interp_frame_paths = sorted(glob.glob(os.path.join(save_folder, '*.png')))
        video_list = []
        for f in interp_frame_paths:
            with Image.open(f) as img:
                img.load()
                video_list.append(img)

        # if saving PNG, also save interpolated frames
        if "PNG" in params.format:
            save_interp_path = f"{p.outpath_samples}/{self.working_dir}/interp"
            os.makedirs(save_interp_path, exist_ok=True)
            shutil.move(save_folder, save_interp_path)

        # remove tmp folder
        try:
            shutil.rmtree(tmp_folder)
        except OSError as e:
            print(f"Error: {e}")

        return video_list

    def _save(
            self,
            params: HotshotXLParams,
            video_list: list,
            video_path_prefix: str,
            res: Processed,
            index: int,
    ):
        video_paths = []
        video_array = [np.array(v) for v in video_list]
        if "GIF" in params.format:
            video_path_gif = video_path_prefix + "gif"
            video_paths.append(video_path_gif)
            if shared.opts.data.get("hotshot_xl_optimize_gif_palette", False):
                try:
                    import av
                except ImportError:
                    from launch import run_pip
                    run_pip(
                        "install imageio[pyav]",
                        "hotshot_xl GIF palette optimization requirement: imageio[pyav]",
                    )
                imageio.imwrite(
                    video_path_gif, video_array, plugin='pyav', fps=params.fps,
                    codec='gif', out_pixel_format='pal8',
                    filter_graph=(
                        {
                            "split": ("split", ""),
                            "palgen": ("palettegen", ""),
                            "paluse": ("paletteuse", ""),
                            "scale": ("scale", f"{video_list[0].width}:{video_list[0].height}")
                        },
                        [
                            ("video_in", "scale", 0, 0),
                            ("scale", "split", 0, 0),
                            ("split", "palgen", 1, 0),
                            ("split", "paluse", 0, 0),
                            ("palgen", "paluse", 0, 1),
                            ("paluse", "video_out", 0, 0),
                        ]
                    )
                )
            else:
                imageio.imwrite(
                    video_path_gif,
                    video_array,
                    duration=(1000 / params.fps),
                    loop=params.loop_number,
                )
            if shared.opts.data.get("hotshot_xl_optimize_gif_gifsicle", False):
                self._optimize_gif(video_path_gif)
        if "MP4" in params.format:
            video_path_mp4 = video_path_prefix + "mp4"
            video_paths.append(video_path_mp4)
            imageio.imwrite(video_path_mp4, video_array, fps=params.fps, codec="h264")

        if "TXT" in params.format and res.images[index].info is not None:
            video_path_txt = video_path_prefix + "txt"
            self._save_txt(params, video_path_txt, res, index)
        return video_paths

    def _optimize_gif(self, video_path: str):
        try:
            import pygifsicle
        except ImportError:
            from launch import run_pip

            run_pip(
                "install pygifsicle",
                "hotshot-xl GIF optimization requirement: pygifsicle",
            )
            import pygifsicle
        finally:
            try:
                pygifsicle.optimize(video_path)
            except FileNotFoundError:
                logger.warn(
                    "gifsicle not found, required for optimized GIFs, try: apt install gifsicle"
                )

    def _save_txt(
            self, params: HotshotXLParams, video_path: str, res: Processed, i: int
    ):
        res.images[i].info["temporal_layers"] = params.model
        res.images[i].info["video_length"] = params.video_length
        res.images[i].info["fps"] = params.fps
        res.images[i].info["loop_number"] = params.loop_number
        with open(video_path, "w", encoding="utf8") as file:
            file.write(f"{res.images[i].info}\n")

    def _encode_video_to_b64(self, paths):
        videos = []
        for v_path in paths:
            with open(v_path, "rb") as video_file:
                encoded_video = base64.b64encode(video_file.read())
            videos.append(encoded_video.decode("utf-8"))
        return videos

scripts/hotshot_xl_output.py
hotshotco/Hotshot-XL-Automatic1111
import os
import cv2
import gradio as gr

import packages

class ToolButton(gr.Button, gr.components.FormComponent):
    """Small button with single emoji as text, fits inside gradio forms"""

    def __init__(self, **kwargs):
        super().__init__(variant="tool", **kwargs)


    def get_block_name(self):
        return "button"

class HotshotXLParams:

    def __init__(
            self,
            model="hsxl_temporal_layers.f16.safetensors",
            enable=False,
            video_length=8,
            fps=8,
            loop_number=0,
            batch_size=8,
            stride=1,
            overlap=-1,
            format=["GIF"],
            #interp='Off',
            #interp_x=10,
            reverse=[],
            original_size_width=1920,
            original_size_height=1080,
            target_size_width=512,
            target_size_height=512,

            negative_original_size_width=512,
            negative_original_size_height=512,
            negative_target_size_width=64,
            negative_target_size_height=64,
    ):
        self.model = model
        self.enable = enable
        self.video_length = video_length
        self.fps = fps
        self.loop_number = loop_number
        self.batch_size = batch_size
        self.stride = stride
        self.overlap = overlap
        self.format = format
        #self.interp = interp
        #self.interp_x = interp_x
        self.reverse = reverse

        self.original_size_width = original_size_width
        self.original_size_height = original_size_height
        self.target_size_width = target_size_width
        self.target_size_height = target_size_height

        self.negative_original_size_width = negative_original_size_width
        self.negative_original_size_height = negative_original_size_height
        self.negative_target_size_width = negative_target_size_width
        self.negative_target_size_height = negative_target_size_height

    def get_list(self, is_img2img: bool):
        list_var = list(vars(self).values())

        return list_var


    def _check(self):
        assert (
            self.video_length >= 0 and self.fps > 0
        ), "Video length and FPS should be positive."
        assert not set(["GIF", "PNG"]).isdisjoint(
            self.format
        ), "At least one saving format should be selected."

    def set_p(self, p):
        self._check()
        if self.video_length < self.batch_size:
            p.batch_size = self.batch_size
        else:
            p.batch_size = self.video_length
        if self.video_length == 0:
            self.video_length = p.batch_size
            self.video_default = True
        else:
            self.video_default = False
        if self.overlap == -1:
            self.overlap = self.batch_size // 4
        if "PNG" not in self.format:
            p.do_not_save_samples = True


class HotshotXLUiGroup:
    txt2img_submit_button = None
    img2img_submit_button = None

    def __init__(self):
        self.params = HotshotXLParams()

    def render(self, is_img2img: bool, model_dir: str):

        if is_img2img:
            return gr.State()

        if not os.path.isdir(model_dir):
            os.mkdir(model_dir)
        elemid_prefix = "img2img-hsxl-" if is_img2img else "txt2img-hsxl-"
        model_list = [f for f in os.listdir(model_dir) if f != ".gitkeep"]

        def padding():
            with gr.Row():
                gr.Markdown('<div style="width:100%; margin-bottom:1em;"></div>')

        with gr.Accordion("Hotshot-XL", open=False):

            further_setup_needed = False

            if not packages.is_diffusers_installed():
                further_setup_needed = True
                with gr.Row():
                    gr.Markdown(f"""This package relies on diffusers=={packages.min_diffusers_version}. Please make sure it is installed!
                    pip install diffusers=={packages.min_diffusers_version}""")

            if len(model_list) == 0:
                further_setup_needed = True
                with gr.Row():
                    gr.Markdown("No models found!")
                with gr.Row():
                    gr.Markdown(f"Please Install models to '{model_dir}'")
                with gr.Row():
                    gr.Markdown(f"""Download the weights from <a href='https://huggingface.co/hotshotco/Hotshot-XL/blob/main/hsxl_temporal_layers.f16.safetensors'>here</a>""")

            if further_setup_needed:
                with gr.Row():
                    gr.Markdown(f"After completing the above, please reload your UI")
                return gr.State()

            with gr.Row():
                gr.Markdown("""This model was trained at aspects close to 512x512 so we recommend using the <a href='https://huggingface.co/hotshotco/SDXL-512/blob/main/hsxl_base_1.0.f16.safetensors'>hotshotco/SDXL-512</a> model which was fine tuned off SDXL.
                For more info about this project you can read more at the <a href='https://github.com/hotshotco/Hotshot-XL'>Github repository.</a>""")

            padding()

            with gr.Row():

                self.params.enable = gr.Checkbox(
                    value=self.params.enable, label="Enabled",
                    elem_id=f"{elemid_prefix}enable"
                )

                def refresh_models(*inputs):
                    new_model_list = [
                        f for f in os.listdir(model_dir) if f != ".gitkeep"
                    ]
                    dd = inputs[0]
                    if dd in new_model_list:
                        selected = dd
                    elif len(new_model_list) > 0:
                        selected = new_model_list[0]
                    else:
                        selected = None
                    return gr.Dropdown.update(choices=new_model_list, value=selected)

                self.params.model = gr.Dropdown(
                    choices=model_list,
                    value=(self.params.model if self.params.model in model_list else None),
                    label="Model",
                    type="value",
                    tooltip="Choose which temporal layers will be injected into the generation process.",
                    elem_id=f"{elemid_prefix}temporal-layers",
                )

                refresh_model = ToolButton(value="\U0001f504")
                refresh_model.click(
                    refresh_models, self.params.model, self.params.model
                )

            padding()

            with gr.Row():

                self.params.video_length = gr.Number(
                    minimum=0,
                    value=self.params.video_length,
                    label="Total Frames",
                    precision=0,
                    tooltip="Total length of video in frames.",
                    elem_id=f"{elemid_prefix}video-length",
                )
                self.params.fps = gr.Number(
                    value=self.params.fps, label="FPS", precision=0,
                    tooltip="How many frames per second the gif will run.",
                    elem_id=f"{elemid_prefix}fps"
                )
                self.params.loop_number = gr.Number(
                    minimum=0,
                    value=self.params.loop_number,
                    label="Loop Count",
                    precision=0,
                    tooltip="How many times the animation will loop, a value of 0 will loop forever.",
                    elem_id=f"{elemid_prefix}loop-number",
                )

            padding()

            with gr.Row():
                # self.params.closed_loop = gr.Checkbox(
                #     value=self.params.closed_loop,
                #     label="Closed loop",
                #     tooltip="If enabled, will try to make the last frame the same as the first frame.",
                #     elem_id=f"{elemid_prefix}closed-loop",
                # )

                self.params.batch_size = gr.Slider(
                    minimum=1,
                    maximum=32,
                    value=self.params.batch_size,
                    label="Context batch size",
                    step=1,
                    precision=0,
                    elem_id=f"{elemid_prefix}batch-size",
                )
                self.params.stride = gr.Number(
                    minimum=1,
                    value=self.params.stride,
                    label="Stride",
                    precision=0,
                    tooltip="",
                    elem_id=f"{elemid_prefix}stride",
                )
                self.params.overlap = gr.Number(
                    minimum=-1,
                    value=self.params.overlap,
                    label="Overlap",
                    precision=0,
                    tooltip="Number of frames to overlap in context.",
                    elem_id=f"{elemid_prefix}overlap",
                )

            padding()

            with gr.Row():
                self.params.format = gr.CheckboxGroup(
                    choices=["GIF", "PNG"],
                    label="Output Format",
                    type="value",
                    tooltip="Which formats the animation should be saved in",
                    elem_id=f"{elemid_prefix}save-format",
                    value=self.params.format,
                )
                self.params.reverse = gr.CheckboxGroup(
                    choices=["Enabled", "Remove head", "Remove tail"],
                    label="Ping Pong Effect",
                    type="index",
                    tooltip="Reverse the resulting animation, remove the first and/or last frame from duplication.",
                    elem_id=f"{elemid_prefix}reverse",
                    value=self.params.reverse
                )

            padding()

            with gr.Row():
                gr.Markdown("SDXL Conditioning")

            with gr.Row():
                with gr.Accordion("Positive"):
                    with gr.Row():
                        self.params.original_size_width = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.original_size_width,
                            label="Original Width",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}og-width",
                        )
                        self.params.target_size_width = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.target_size_width,
                            label="Target Width",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}tgt-width",
                        )

                    padding()

                    with gr.Row():
                        self.params.original_size_height = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.original_size_height,
                            label="Original Height",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}og-height",
                        )
                        self.params.target_size_height = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.target_size_height,
                            label="Target Height",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}tgt-height",
                        )

            with gr.Row():
                with gr.Accordion("Negative"):
                    with gr.Row():
                        self.params.negative_original_size_width = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.negative_original_size_width,
                            label="Original Width",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}ng-og-width",
                        )
                        self.params.negative_target_size_width = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.negative_target_size_width,
                            label="Target Width",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}ng-tgt-width",
                        )

                    padding()

                    with gr.Row():
                        self.params.negative_original_size_height = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.negative_original_size_height,
                            label="Original Height",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}ng-og-height",
                        )
                        self.params.negative_target_size_height = gr.Slider(
                            minimum=64,
                            maximum=4096,
                            value=self.params.negative_target_size_height,
                            label="Target Height",
                            step=8,
                            precision=0,
                            elem_id=f"{elemid_prefix}ng-tgt-height",
                        )

            # with gr.Row():
            #     self.params.interp = gr.Radio(
            #         choices=["Off", "FILM"],
            #         label="Frame Interpolation Mode",
            #         tooltip="Interpolate between frames with Deforum's FILM implementation. Requires Deforum extension.",
            #         elem_id=f"{elemid_prefix}interp-choice",
            #         value=self.params.interp
            #     )
            #     self.params.interp_x = gr.Number(
            #         value=self.params.interp_x, label="Interp X", precision=0,
            #         tooltip="Replace each input frame with X interpolated output frames.",
            #         elem_id=f"{elemid_prefix}interp-x"
            #     )
            # self.params.video_source = gr.Video(
            #     value=self.params.video_source,
            #     label="Video source",
            # )

            def update_fps(video_source):
                if video_source is not None and video_source != '':
                    cap = cv2.VideoCapture(video_source)
                    fps = int(cap.get(cv2.CAP_PROP_FPS))
                    cap.release()
                    return fps
                else:
                    return int(self.params.fps.value)

            #self.params.video_source.change(update_fps, inputs=self.params.video_source, outputs=self.params.fps)
            padding()
            with gr.Row():
                unload = gr.Button(
                    value="Unload Model"
                )
                unload.click(fn=self.unload_model)
        return self.register_unit(is_img2img)

    def unload_model(self):
        from scripts.hotshot_xl_model_controller import model_controller
        model_controller.unload()

    def register_unit(self, is_img2img: bool):

        unit = gr.State(value=HotshotXLParams)
        button = HotshotXLUiGroup.img2img_submit_button if is_img2img else HotshotXLUiGroup.txt2img_submit_button
        if button:
            button.click(
                fn=HotshotXLParams,
                inputs=self.params.get_list(is_img2img),
                outputs=unit,
                queue=False,
            )

        return unit


    @staticmethod
    def on_after_component(component, **_kwargs):
        elem_id = getattr(component, "elem_id", None)

        if elem_id == "txt2img_generate":
            HotshotXLUiGroup.txt2img_submit_button = component
            return

        if elem_id == "img2img_generate":
            HotshotXLUiGroup.img2img_submit_button = component
            return
scripts/hotshot_xl_ui.py
hotshotco/Hotshot-XL-Automatic1111