`.typings/mflux/callbacks/callback.pyi`

48 lines
"""
This type stub file was generated by pyright.
"""

import mlx.core as mx
import PIL.Image
import tqdm
from typing import Protocol
from mflux.models.common.config.config import Config

class BeforeLoopCallback(Protocol):
    def call_before_loop(
        self,
        seed: int,
        prompt: str,
        latents: mx.array,
        config: Config,
        canny_image: PIL.Image.Image | None = ...,
        depth_image: PIL.Image.Image | None = ...,
    ) -> None: ...

class InLoopCallback(Protocol):
    def call_in_loop(
        self,
        t: int,
        seed: int,
        prompt: str,
        latents: mx.array,
        config: Config,
        time_steps: tqdm,
    ) -> None: ...

class AfterLoopCallback(Protocol):
    def call_after_loop(
        self, seed: int, prompt: str, latents: mx.array, config: Config
    ) -> None: ...

class InterruptCallback(Protocol):
    def call_interrupt(
        self,
        t: int,
        seed: int,
        prompt: str,
        latents: mx.array,
        config: Config,
        time_steps: tqdm,
    ) -> None: ...