`.typings/mflux/models/flux/variants/txt2img/flux.pyi`

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

from pathlib import Path
from mlx import nn
from typing import Any
from mflux.models.common.config.model_config import ModelConfig
from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import (
    CLIPEncoder,
)
from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.flux.model.flux_transformer.transformer import Transformer
from mflux.models.flux.model.flux_vae.vae import VAE
from mflux.utils.generated_image import GeneratedImage

class Flux1(nn.Module):
    vae: VAE
    transformer: Transformer
    t5_text_encoder: T5Encoder
    clip_text_encoder: CLIPEncoder
    bits: int | None
    lora_paths: list[str] | None
    lora_scales: list[float] | None
    prompt_cache: dict[str, Any]
    tokenizers: dict[str, Any]

    def __init__(
        self,
        quantize: int | None = ...,
        model_path: str | None = ...,
        lora_paths: list[str] | None = ...,
        lora_scales: list[float] | None = ...,
        model_config: ModelConfig = ...,
    ) -> None: ...
    def generate_image(
        self,
        seed: int,
        prompt: str,
        num_inference_steps: int = ...,
        height: int = ...,
        width: int = ...,
        guidance: float = ...,
        image_path: Path | str | None = ...,
        image_strength: float | None = ...,
        scheduler: str = ...,
        negative_prompt: str | None = ...,
    ) -> GeneratedImage: ...
    @staticmethod
    def from_name(model_name: str, quantize: int | None = ...) -> Flux1: ...
    def save_model(self, base_path: str) -> None: ...
    def freeze(self, **kwargs):  # -> None:
        ...