#!/usr/bin/env python3
"""Text-to-image with Bria's FIBO 2.0 Edit Turbo on Apple silicon: fibo2_edit_turbo_t2i_mlx.py with mflux's FLUX.2 VAE.

Everything except the VAE decode comes from fibo2_edit_turbo_t2i_mlx.py, so keep the two files together: the transformer,
the text encoder, the captions, the options and the output are the same. The final latent is decoded by mflux's FLUX.2
VAE (https://github.com/mflux-community/mflux), loaded from the model's own vae/ folder, for those who already work with
mflux. FIBO 2.0 itself is not an mflux model; only its VAE, FLUX.2's, is.

Tested with mflux 0.22.0. The VAE classes used here are mflux internals, not a public API, and may move between mflux
releases; fibo2_edit_turbo_t2i_mlx.py does the same decode without mflux.

    python fibo2_edit_turbo_t2i_mflux.py "a red vintage bicycle leaning against a whitewashed wall on a narrow street"
    python fibo2_edit_turbo_t2i_mflux.py --json assets/lake-pier.json --seed 4242

Requirements: requirements-mflux.txt (fibo2_edit_turbo_t2i_mlx.py's, plus mflux).
"""
import dataclasses
import os
import sys
from pathlib import Path

import mlx.core as mx
import numpy as np
from mlx.utils import tree_flatten
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import fibo2_edit_turbo_t2i_mlx as fibo2_mlx  # noqa: E402


class MfluxVae:
    """mflux's Flux2VAE with the model's vae/ weights: decode(latent) -> PIL, as fibo2_edit_turbo_t2i_mlx.VaeDecoder."""

    def __init__(self, folder):
        from mflux.models.common.weights.loading.weight_applier import WeightApplier
        from mflux.models.common.weights.loading.weight_loader import WeightLoader
        from mflux.models.flux2.model.flux2_vae.vae import Flux2VAE
        from mflux.models.flux2.weights.flux2_weight_definition import Flux2KleinWeightDefinition
        path = Path(folder)
        comp = next(c for c in Flux2KleinWeightDefinition.get_components() if c.name == "vae")
        comp = dataclasses.replace(comp, hf_subdir=path.name)

        class VaeDefinition(Flux2KleinWeightDefinition):
            @staticmethod
            def get_components():
                return [comp]
        self.vae = Flux2VAE()
        weights = WeightLoader.load_single_local(component=comp, root_path=path.parent)
        got = dict(tree_flatten(weights.components["vae"]))
        missing = [k for k, _ in tree_flatten(self.vae.parameters()) if k not in got]
        if missing:            # mflux applies weights leniently: a missing one would stay at its random start
            raise ValueError(f"{folder}: {len(missing)} VAE parameters missing for mflux's Flux2VAE ({missing[:3]})")
        WeightApplier.apply_and_quantize(weights=weights, quantize_arg=None, weight_definition=VaeDefinition,
                                         models={"vae": self.vae})
        bn = self.vae.bn
        self.mean = bn.running_mean.reshape(1, -1, 1, 1)
        self.std = mx.sqrt(bn.running_var.reshape(1, -1, 1, 1) + bn.eps)

    def decode(self, latents):
        B, C, H, W = latents.shape
        grid = latents.astype(mx.float32).reshape(B, C, H // 2, 2, W // 2, 2).transpose(0, 1, 3, 5, 2, 4)
        grid = grid.reshape(B, C * 4, H // 2, W // 2) * self.std + self.mean
        x = self.vae.decode(self.vae._unpatchify_latents(grid))
        rgb = mx.clip(x[0].transpose(1, 2, 0).astype(mx.float32) / 2 + 0.5, 0, 1) * 255
        mx.eval(rgb)
        return Image.fromarray(np.array(rgb).round().astype(np.uint8))


def load_decoder(folder, dtype):
    return MfluxVae(os.path.join(folder, "vae"))


if __name__ == "__main__":
    sys.exit(fibo2_mlx.main(load_decoder=load_decoder,
                            description="Text-to-image with FIBO 2.0 Edit Turbo on Apple silicon (MLX, mflux's VAE)"))