`.typings/mflux/models/qwen/variants/edit/qwen_image_edit.pyi`

56 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 Config
from mflux.models.common.config.model_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 QwenImageEdit(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,
        image_paths: list[str],
        num_inference_steps: int = ...,
        height: int | None = ...,
        width: int | None = ...,
        guidance: float = ...,
        image_path: Path | str | None = ...,
        scheduler: str = ...,
        negative_prompt: str | None = ...,
    ) -> GeneratedImage: ...
    def _encode_prompts_with_images(
        self,
        prompt: str,
        negative_prompt: str,
        image_paths: list[str],
        config: Config,
        vl_width: int | None,
        vl_height: int | None,
    ) -> tuple[mx.array, mx.array, mx.array, mx.array]: ...