`.typings/mflux/models/qwen/variants/txt2img/qwen_image.pyi`

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

import mlx.core as mx
from pathlib import Path
from mlx import nn
from typing import Any
from mflux.models.common.config import ModelConfig
from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE
from mflux.utils.generated_image import GeneratedImage

class QwenImage(nn.Module):
    vae: QwenVAE
    transformer: QwenTransformer
    text_encoder: QwenTextEncoder
    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: ...
    def save_model(self, base_path: str) -> None: ...
    @staticmethod
    def compute_guided_noise(
        noise: mx.array, noise_negative: mx.array, guidance: float
    ) -> mx.array: ...