pipeline_output.py 596 Bytes
Newer Older
Aryan V S's avatar
Aryan V S committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from dataclasses import dataclass
from typing import List, Union

import numpy as np
import PIL.Image
import torch

from ...utils import BaseOutput


@dataclass
class AnimateDiffPipelineOutput(BaseOutput):
    r"""
    Output class for AnimateDiff pipelines.

    Args:
        frames (`List[List[PIL.Image.Image]]` or `torch.Tensor` or `np.ndarray`):
            List of PIL Images of length `batch_size` or torch.Tensor or np.ndarray of shape
            `(batch_size, num_frames, height, width, num_channels)`.
    """

    frames: Union[List[List[PIL.Image.Image]], torch.Tensor, np.ndarray]