# Quanto optimization, unique
import os
import torch
import torch._dynamo
import gc
import json
import transformers

from huggingface_hub.constants import HF_HUB_CACHE
from transformers import T5EncoderModel
import diffusers
from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
from torch import Generator
from diffusers import FluxTransformer2DModel, DiffusionPipeline

from PIL.Image import Image
from diffusers import AutoencoderTiny
from pipelines.models import TextToImageRequest
from optimum.quanto import requantize as optimum_quant

try:
    from huggingface_hub import hf_hub_download
except:
    pass


torch._dynamo.config.suppress_errors = True
os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "True"

ckpt_main = "black-forest-labs/FLUX.1-schnell"
revision_main = "741f7c3ce8b383c54771c7003378a50191e9efe9"
Pipeline = None
apply_transformer_tag = 1

import torch
import gc
import os
import json
import transformers


def convert_transformer_to_int8(repo_path):
    with open("transformer_int8.json", "r") as f:
        quantization_map = json.load(f)
    with torch.device("meta"):
        transformer_config_path = os.path.join(repo_path, "config.json")
        transformer = diffusers.FluxTransformer2DModel.from_config(transformer_config_path).to(torch.bfloat16)

    state_dict = hf_hub_download(repo_path, "diffusion_pytorch_models.safetensors")

    optimum_quant(transformer, state_dict, quantization_map, device=torch.device("cuda"))
    return transformer


def load_pipeline() -> Pipeline:

    original_vae = AutoencoderTiny.from_pretrained("RichardWilliam/XULF_Vae",
                    revision="3ee225c539465c27adadec45c6e8af50a7397b7d",
                    torch_dtype=torch.bfloat16)
    

    text_encoder_2 =  T5EncoderModel.from_pretrained("RichardWilliam/XULF_T5_bf16",
                revision = "63a3d9ef7b586655600ac9bd4e4747d038237761", 
                torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)

    trans_path = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_Transfomer/snapshots/6860c51af40329808f270e159a0d018559a1204f")
    pre_quanted_trans = FluxTransformer2DModel.from_pretrained(trans_path,
                        torch_dtype=torch.bfloat16, 
                        use_safetensors=False).to(memory_format=torch.channels_last)
    transformer = pre_quanted_trans

    pipeline = DiffusionPipeline.from_pretrained(ckpt_main, 
                        revision=revision_main, 
                        vae=original_vae, 
                        transformer=transformer, 
                        text_encoder_2=text_encoder_2, 
                        torch_dtype=torch.bfloat16)
    pipeline.to("cuda")
    try:
        pipeline.enable_int8()
        pipeline.transformer = convert_transformer_to_int8(trans_path)
    except:
        print("Use origin pipeline")

    for warm_up_prompt in range(3):
        pipeline(prompt="puffer, cutie, buttinsky, prototrophic, betulinamaric, quintet, tunesome, decaspermous", 
                        width=1024,
                        height=1024,
                        guidance_scale=0.0,
                        num_inference_steps=4,
                        max_sequence_length=256)
    return pipeline

@torch.no_grad()
def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:

    gc.collect()
    torch.cuda.empty_cache()

    generator = Generator(pipeline.device).manual_seed(request.seed)

    return pipeline(
        request.prompt,
        generator=generator,
        guidance_scale=0.0,
        num_inference_steps=4,
        max_sequence_length=256,
        height=request.height,
        width=request.width,
    ).images[0]