stimulus_synthesis/generators/base.py

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any


class TextToImageGenerator(ABC):
    @abstractmethod
    def generate(self, prompts: list[str], *, generator: Any | None = None, **kwargs) -> list[Any]:
        ...


class ImageToVideoGenerator(ABC):
    @abstractmethod
    def generate(self, image: Any, prompt: str, *, generator: Any | None = None, **kwargs) -> Any:
        ...

    def generate_batch(self, images: list[Any], prompts: list[str], *, generators: list[Any] | None = None, **kwargs) -> list[Any]:
        generators = generators or [None] * len(images)
        return [
            self.generate(image, prompt, generator=gen, **kwargs)
            for image, prompt, gen in zip(images, prompts, generators)
        ]