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

import mlx.core as mx
import PIL.Image
import tqdm
from typing import TYPE_CHECKING
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.models.common.config.config import Config

if TYPE_CHECKING: ...

class GenerationContext:
    def __init__(
        self, registry: CallbackRegistry, seed: int, prompt: str, config: Config
    ) -> None: ...
    def before_loop(
        self,
        latents: mx.array,
        *,
        canny_image: PIL.Image.Image | None = ...,
        depth_image: PIL.Image.Image | None = ...,
    ) -> None: ...
    def in_loop(self, t: int, latents: mx.array, time_steps: tqdm = ...) -> None: ...
    def after_loop(self, latents: mx.array) -> None: ...
    def interruption(
        self, t: int, latents: mx.array, time_steps: tqdm = ...
    ) -> None: ...
