`.typings/mflux/models/flux/variants/concept_attention/attention_data.pyi`

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

import mlx.core as mx
import PIL.Image
from dataclasses import dataclass
from pathlib import Path
from typing import List
from mflux.models.flux.variants.concept_attention.joint_transformer_block_concept import (
    LayerAttentionData,
)

@dataclass
class TimestepAttentionData:
    t: int
    attention_information: List[LayerAttentionData]
    def stack_img_attentions(self) -> mx.array: ...
    def stack_concept_attentions(self) -> mx.array: ...

class GenerationAttentionData:
    def __init__(self) -> None: ...
    def append(self, timestep_attention: TimestepAttentionData):  # -> None:
        ...
    def stack_all_img_attentions(self) -> mx.array: ...
    def stack_all_concept_attentions(self) -> mx.array: ...

@dataclass
class ConceptHeatmap:
    concept: str
    image: PIL.Image.Image
    layer_indices: List[int]
    timesteps: List[int]
    height: int
    width: int
    def save(
        self, path: str | Path, export_json_metadata: bool = ..., overwrite: bool = ...
    ) -> None: ...
    def get_metadata(self) -> dict: ...