generation.py

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum

import torch


class AdvanceStatus(str, Enum):
    WAIT_INPUT = "WAIT_INPUT"
    PRODUCED = "PRODUCED"
    END_AUDIO = "END_AUDIO"
    FORCED_STOP = "FORCED_STOP"


@dataclass
class GenerationResult:
    frames: list[torch.Tensor]
    status: AdvanceStatus
    actions: list[int]
    consumed_text_tokens: int
    text_end_consumed: bool