xinference 0.15.4__py3-none-any.whl → 0.16.1__py3-none-any.whl

xinference/thirdparty/mlx/flux/model.py

@@ -0,0 +1,134 @@
+# Copyright © 2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    LastLayer,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class Flux(nn.Module):
+    def __init__(self, params: FluxParams):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(
+                f"Got {params.axes_dim} but expected positional dim {pe_dim}"
+            )
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(
+            dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim
+        )
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+            if params.guidance_embed
+            else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = [
+            DoubleStreamBlock(
+                self.hidden_size,
+                self.num_heads,
+                mlp_ratio=params.mlp_ratio,
+                qkv_bias=params.qkv_bias,
+            )
+            for _ in range(params.depth)
+        ]
+
+        self.single_blocks = [
+            SingleStreamBlock(
+                self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio
+            )
+            for _ in range(params.depth_single_blocks)
+        ]
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def sanitize(self, weights):
+        new_weights = {}
+        for k, w in weights.items():
+            if k.endswith(".scale"):
+                k = k[:-6] + ".weight"
+            for seq in ["img_mlp", "txt_mlp", "adaLN_modulation"]:
+                if f".{seq}." in k:
+                    k = k.replace(f".{seq}.", f".{seq}.layers.")
+                    break
+            new_weights[k] = w
+        return new_weights
+
+    def __call__(
+        self,
+        img: mx.array,
+        img_ids: mx.array,
+        txt: mx.array,
+        txt_ids: mx.array,
+        timesteps: mx.array,
+        y: mx.array,
+        guidance: Optional[mx.array] = None,
+    ) -> mx.array:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError(
+                    "Didn't get guidance strength for guidance distilled model."
+                )
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = mx.concatenate([txt_ids, img_ids], axis=1)
+        pe = self.pe_embedder(ids).astype(img.dtype)
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+        img = mx.concatenate([txt, img], axis=1)
+        for block in self.single_blocks:
+            img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)
+
+        return img

xinference/thirdparty/mlx/flux/sampler.py

@@ -0,0 +1,56 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from functools import lru_cache
+
+import mlx.core as mx
+
+
+class FluxSampler:
+    def __init__(self, name: str, base_shift: float = 0.5, max_shift: float = 1.5):
+        self._base_shift = base_shift
+        self._max_shift = max_shift
+        self._schnell = "schnell" in name
+
+    def _time_shift(self, x, t):
+        x1, x2 = 256, 4096
+        t1, t2 = self._base_shift, self._max_shift
+        exp_mu = math.exp((x - x1) * (t2 - t1) / (x2 - x1) + t1)
+        t = exp_mu / (exp_mu + (1 / t - 1))
+        return t
+
+    @lru_cache
+    def timesteps(
+        self, num_steps, image_sequence_length, start: float = 1, stop: float = 0
+    ):
+        t = mx.linspace(start, stop, num_steps + 1)
+
+        if self._schnell:
+            t = self._time_shift(image_sequence_length, t)
+
+        return t.tolist()
+
+    def random_timesteps(self, B, L, dtype=mx.float32, key=None):
+        if self._schnell:
+            # TODO: Should we upweigh 1 and 0.75?
+            t = mx.random.randint(1, 5, shape=(B,), key=key)
+            t = t.astype(dtype) / 4
+        else:
+            t = mx.random.uniform(shape=(B,), dtype=dtype, key=key)
+            t = self._time_shift(L, t)
+
+        return t
+
+    def sample_prior(self, shape, dtype=mx.float32, key=None):
+        return mx.random.normal(shape, dtype=dtype, key=key)
+
+    def add_noise(self, x, t, noise=None, key=None):
+        noise = (
+            noise
+            if noise is not None
+            else mx.random.normal(x.shape, dtype=x.dtype, key=key)
+        )
+        return x * (1 - t) + t * noise
+
+    def step(self, pred, x_t, t, t_prev):
+        return x_t + (t_prev - t) * pred

xinference/thirdparty/mlx/flux/t5.py

@@ -0,0 +1,244 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+_SHARED_REPLACEMENT_PATTERNS = [
+    (".block.", ".layers."),
+    (".k.", ".key_proj."),
+    (".o.", ".out_proj."),
+    (".q.", ".query_proj."),
+    (".v.", ".value_proj."),
+    ("shared.", "wte."),
+    ("lm_head.", "lm_head.linear."),
+    (".layer.0.layer_norm.", ".ln1."),
+    (".layer.1.layer_norm.", ".ln2."),
+    (".layer.2.layer_norm.", ".ln3."),
+    (".final_layer_norm.", ".ln."),
+    (
+        "layers.0.layer.0.SelfAttention.relative_attention_bias.",
+        "relative_attention_bias.embeddings.",
+    ),
+]
+
+_ENCODER_REPLACEMENT_PATTERNS = [
+    (".layer.0.SelfAttention.", ".attention."),
+    (".layer.1.DenseReluDense.", ".dense."),
+]
+
+
+@dataclass
+class T5Config:
+    vocab_size: int
+    num_layers: int
+    num_heads: int
+    relative_attention_num_buckets: int
+    d_kv: int
+    d_model: int
+    feed_forward_proj: str
+    tie_word_embeddings: bool
+
+    d_ff: Optional[int] = None
+    num_decoder_layers: Optional[int] = None
+    relative_attention_max_distance: int = 128
+    layer_norm_epsilon: float = 1e-6
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(
+            vocab_size=config["vocab_size"],
+            num_layers=config["num_layers"],
+            num_heads=config["num_heads"],
+            relative_attention_num_buckets=config["relative_attention_num_buckets"],
+            d_kv=config["d_kv"],
+            d_model=config["d_model"],
+            feed_forward_proj=config["feed_forward_proj"],
+            tie_word_embeddings=config["tie_word_embeddings"],
+            d_ff=config.get("d_ff", 4 * config["d_model"]),
+            num_decoder_layers=config.get("num_decoder_layers", config["num_layers"]),
+            relative_attention_max_distance=config.get(
+                "relative_attention_max_distance", 128
+            ),
+            layer_norm_epsilon=config.get("layer_norm_epsilon", 1e-6),
+        )
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, config: T5Config, bidirectional: bool):
+        self.bidirectional = bidirectional
+        self.num_buckets = config.relative_attention_num_buckets
+        self.max_distance = config.relative_attention_max_distance
+        self.n_heads = config.num_heads
+        self.embeddings = nn.Embedding(self.num_buckets, self.n_heads)
+
+    @staticmethod
+    def _relative_position_bucket(rpos, bidirectional, num_buckets, max_distance):
+        num_buckets = num_buckets // 2 if bidirectional else num_buckets
+        max_exact = num_buckets // 2
+
+        abspos = rpos.abs()
+        is_small = abspos < max_exact
+
+        scale = (num_buckets - max_exact) / math.log(max_distance / max_exact)
+        buckets_large = (mx.log(abspos / max_exact) * scale).astype(mx.int16)
+        buckets_large = mx.minimum(max_exact + buckets_large, num_buckets - 1)
+
+        buckets = mx.where(is_small, abspos, buckets_large)
+        if bidirectional:
+            buckets = buckets + (rpos > 0) * num_buckets
+        else:
+            buckets = buckets * (rpos < 0)
+
+        return buckets
+
+    def __call__(self, query_length: int, key_length: int, offset: int = 0):
+        """Compute binned relative position bias"""
+        context_position = mx.arange(offset, query_length)[:, None]
+        memory_position = mx.arange(key_length)[None, :]
+
+        # shape (query_length, key_length)
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,
+            bidirectional=self.bidirectional,
+            num_buckets=self.num_buckets,
+            max_distance=self.max_distance,
+        )
+
+        # shape (query_length, key_length, num_heads)
+        values = self.embeddings(relative_position_bucket)
+
+        # shape (num_heads, query_length, key_length)
+        return values.transpose(2, 0, 1)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config: T5Config):
+        super().__init__()
+        inner_dim = config.d_kv * config.num_heads
+        self.num_heads = config.num_heads
+        self.query_proj = nn.Linear(config.d_model, inner_dim, bias=False)
+        self.key_proj = nn.Linear(config.d_model, inner_dim, bias=False)
+        self.value_proj = nn.Linear(config.d_model, inner_dim, bias=False)
+        self.out_proj = nn.Linear(inner_dim, config.d_model, bias=False)
+
+    def __call__(
+        self,
+        queries: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> [mx.array, Tuple[mx.array, mx.array]]:
+        queries = self.query_proj(queries)
+        keys = self.key_proj(keys)
+        values = self.value_proj(values)
+
+        num_heads = self.num_heads
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        queries = queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, S, num_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, S, num_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            key_cache, value_cache = cache
+            keys = mx.concatenate([key_cache, keys], axis=3)
+            values = mx.concatenate([value_cache, values], axis=2)
+
+        values_hat = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=1.0, mask=mask.astype(queries.dtype)
+        )
+        values_hat = values_hat.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.out_proj(values_hat), (keys, values)
+
+
+class DenseActivation(nn.Module):
+    def __init__(self, config: T5Config):
+        super().__init__()
+        mlp_dims = config.d_ff or config.d_model * 4
+        self.gated = config.feed_forward_proj.startswith("gated")
+        if self.gated:
+            self.wi_0 = nn.Linear(config.d_model, mlp_dims, bias=False)
+            self.wi_1 = nn.Linear(config.d_model, mlp_dims, bias=False)
+        else:
+            self.wi = nn.Linear(config.d_model, mlp_dims, bias=False)
+        self.wo = nn.Linear(mlp_dims, config.d_model, bias=False)
+        activation = config.feed_forward_proj.removeprefix("gated-")
+        if activation == "relu":
+            self.act = nn.relu
+        elif activation == "gelu":
+            self.act = nn.gelu
+        elif activation == "silu":
+            self.act = nn.silu
+        else:
+            raise ValueError(f"Unknown activation: {activation}")
+
+    def __call__(self, x):
+        if self.gated:
+            hidden_act = self.act(self.wi_0(x))
+            hidden_linear = self.wi_1(x)
+            x = hidden_act * hidden_linear
+        else:
+            x = self.act(self.wi(x))
+        return self.wo(x)
+
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, config: T5Config):
+        super().__init__()
+        self.attention = MultiHeadAttention(config)
+        self.ln1 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.ln2 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dense = DenseActivation(config)
+
+    def __call__(self, x, mask):
+        y = self.ln1(x)
+        y, _ = self.attention(y, y, y, mask=mask)
+        x = x + y
+
+        y = self.ln2(x)
+        y = self.dense(y)
+        return x + y
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, config: T5Config):
+        super().__init__()
+        self.layers = [
+            TransformerEncoderLayer(config) for i in range(config.num_layers)
+        ]
+        self.ln = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.relative_attention_bias = RelativePositionBias(config, bidirectional=True)
+
+    def __call__(self, x: mx.array):
+        pos_bias = self.relative_attention_bias(x.shape[1], x.shape[1])
+        pos_bias = pos_bias.astype(x.dtype)
+        for layer in self.layers:
+            x = layer(x, mask=pos_bias)
+        return self.ln(x)
+
+
+class T5Encoder(nn.Module):
+    def __init__(self, config: T5Config):
+        self.wte = nn.Embedding(config.vocab_size, config.d_model)
+        self.encoder = TransformerEncoder(config)
+
+    def sanitize(self, weights):
+        new_weights = {}
+        for k, w in weights.items():
+            for old, new in _SHARED_REPLACEMENT_PATTERNS:
+                k = k.replace(old, new)
+            if k.startswith("encoder."):
+                for old, new in _ENCODER_REPLACEMENT_PATTERNS:
+                    k = k.replace(old, new)
+            new_weights[k] = w
+        return new_weights
+
+    def __call__(self, inputs: mx.array):
+        return self.encoder(self.wte(inputs))

xinference/thirdparty/mlx/flux/tokenizers.py

@@ -0,0 +1,185 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import regex
+from sentencepiece import SentencePieceProcessor
+
+
+class CLIPTokenizer:
+    """A simple port of CLIPTokenizer from https://github.com/huggingface/transformers/ ."""
+
+    def __init__(self, bpe_ranks, vocab, max_length=77):
+        self.max_length = max_length
+        self.bpe_ranks = bpe_ranks
+        self.vocab = vocab
+        self.pat = regex.compile(
+            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            regex.IGNORECASE,
+        )
+
+        self._cache = {self.bos: self.bos, self.eos: self.eos}
+
+    @property
+    def bos(self):
+        return "<|startoftext|>"
+
+    @property
+    def bos_token(self):
+        return self.vocab[self.bos]
+
+    @property
+    def eos(self):
+        return "<|endoftext|>"
+
+    @property
+    def eos_token(self):
+        return self.vocab[self.eos]
+
+    def bpe(self, text):
+        if text in self._cache:
+            return self._cache[text]
+
+        unigrams = list(text[:-1]) + [text[-1] + "</w>"]
+        unique_bigrams = set(zip(unigrams, unigrams[1:]))
+
+        if not unique_bigrams:
+            return unigrams
+
+        # In every iteration try to merge the two most likely bigrams. If none
+        # was merged we are done.
+        #
+        # Ported from https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/tokenization_clip.py
+        while unique_bigrams:
+            bigram = min(
+                unique_bigrams, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))
+            )
+            if bigram not in self.bpe_ranks:
+                break
+
+            new_unigrams = []
+            skip = False
+            for a, b in zip(unigrams, unigrams[1:]):
+                if skip:
+                    skip = False
+                    continue
+
+                if (a, b) == bigram:
+                    new_unigrams.append(a + b)
+                    skip = True
+
+                else:
+                    new_unigrams.append(a)
+
+            if not skip:
+                new_unigrams.append(b)
+
+            unigrams = new_unigrams
+            unique_bigrams = set(zip(unigrams, unigrams[1:]))
+
+        self._cache[text] = unigrams
+
+        return unigrams
+
+    def tokenize(self, text, prepend_bos=True, append_eos=True):
+        if isinstance(text, list):
+            return [self.tokenize(t, prepend_bos, append_eos) for t in text]
+
+        # Lower case cleanup and split according to self.pat. Hugging Face does
+        # a much more thorough job here but this should suffice for 95% of
+        # cases.
+        clean_text = regex.sub(r"\s+", " ", text.lower())
+        tokens = regex.findall(self.pat, clean_text)
+
+        # Split the tokens according to the byte-pair merge file
+        bpe_tokens = [ti for t in tokens for ti in self.bpe(t)]
+
+        # Map to token ids and return
+        tokens = [self.vocab[t] for t in bpe_tokens]
+        if prepend_bos:
+            tokens = [self.bos_token] + tokens
+        if append_eos:
+            tokens.append(self.eos_token)
+
+        if len(tokens) > self.max_length:
+            tokens = tokens[: self.max_length]
+            if append_eos:
+                tokens[-1] = self.eos_token
+
+        return tokens
+
+    def encode(self, text):
+        if not isinstance(text, list):
+            return self.encode([text])
+
+        tokens = self.tokenize(text)
+        length = max(len(t) for t in tokens)
+        for t in tokens:
+            t.extend([self.eos_token] * (length - len(t)))
+
+        return mx.array(tokens)
+
+
+class T5Tokenizer:
+    def __init__(self, model_file, max_length=512):
+        self._tokenizer = SentencePieceProcessor(model_file)
+        self.max_length = max_length
+
+    @property
+    def pad(self):
+        try:
+            return self._tokenizer.id_to_piece(self.pad_token)
+        except IndexError:
+            return None
+
+    @property
+    def pad_token(self):
+        return self._tokenizer.pad_id()
+
+    @property
+    def bos(self):
+        try:
+            return self._tokenizer.id_to_piece(self.bos_token)
+        except IndexError:
+            return None
+
+    @property
+    def bos_token(self):
+        return self._tokenizer.bos_id()
+
+    @property
+    def eos(self):
+        try:
+            return self._tokenizer.id_to_piece(self.eos_token)
+        except IndexError:
+            return None
+
+    @property
+    def eos_token(self):
+        return self._tokenizer.eos_id()
+
+    def tokenize(self, text, prepend_bos=True, append_eos=True, pad=True):
+        if isinstance(text, list):
+            return [self.tokenize(t, prepend_bos, append_eos, pad) for t in text]
+
+        tokens = self._tokenizer.encode(text)
+
+        if prepend_bos and self.bos_token >= 0:
+            tokens = [self.bos_token] + tokens
+        if append_eos and self.eos_token >= 0:
+            tokens.append(self.eos_token)
+        if pad and len(tokens) < self.max_length and self.pad_token >= 0:
+            tokens += [self.pad_token] * (self.max_length - len(tokens))
+
+        return tokens
+
+    def encode(self, text, pad=True):
+        if not isinstance(text, list):
+            return self.encode([text], pad=pad)
+
+        pad_token = self.pad_token if self.pad_token >= 0 else 0
+        tokens = self.tokenize(text, pad=pad)
+        length = max(len(t) for t in tokens)
+        for t in tokens:
+            t.extend([pad_token] * (length - len(t)))
+
+        return mx.array(tokens)