import os
import torch
import torch._dynamo
import gc
import bitsandbytes as bnb
from bitsandbytes.nn.modules import Params4bit, QuantState
import json
import transformers
from huggingface_hub.constants import HF_HUB_CACHE
from transformers import T5EncoderModel, T5TokenizerFast

from torch import Generator
from diffusers import FluxTransformer2DModel, DiffusionPipeline

from PIL.Image import Image
from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
from pipelines.models import TextToImageRequest
import json



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

CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
Pipeline = None



def quantized_matrix_multiply(x, weight, bias):
    """Perform matrix multiplication for 4-bit quantized weights."""
    output = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
    return output.to(x)

def copy_quant_state(state, device=None):
    """Create a copy of quantization state for a given device."""
    if state is None:
        return None

    device = device or state.absmax.device
    nested_state = (
        QuantState(
            absmax=state.state2.absmax.to(device),
            shape=state.state2.shape,
            code=state.state2.code.to(device),
            blocksize=state.state2.blocksize,
            quant_type=state.state2.quant_type,
            dtype=state.state2.dtype,
        )
        if state.nested else None
    )

    return QuantState(
        absmax=state.absmax.to(device),
        shape=state.shape,
        code=state.code,
        blocksize=state.blocksize,
        quant_type=state.quant_type,
        dtype=state.dtype,
        offset=state.offset.to(device) if state.nested else None,
        state2=nested_state,
    )

class QuantizedModelParams(Params4bit):
    def to(self, *args, **kwargs):
        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
        if device is not None and device.type == "cuda" and not self.bnb_quantized:
            return self._quantize(device)
        
        updated_params = QuantizedModelParams(
            torch.nn.Parameter.to(self, device=device, dtype=dtype, non_blocking=non_blocking),
            requires_grad=self.requires_grad,
            quant_state=copy_quant_state(self.quant_state, device),
            compress_statistics=False,
            blocksize=64,
            quant_type=self.quant_type,
            quant_storage=self.quant_storage,
            bnb_quantized=self.bnb_quantized,
            module=self.module
        )
        self.module.quant_state = updated_params.quant_state
        self.data = updated_params.data
        self.quant_state = updated_params.quant_state
        return updated_params

class QuantizedLinearLayer(torch.nn.Module):
    def __init__(self, *args, device=None, dtype=None, **kwargs):
        super().__init__()
        self.dummy = torch.nn.Parameter(torch.empty(1, device=device, dtype=dtype))
        self.weight = None
        self.quant_state = None
        self.bias = None
        self.quant_type = 'nf4'

    def forward(self, x):
        self.weight.quant_state = self.quant_state
        if self.bias is not None and self.bias.dtype != x.dtype:
            self.bias.data = self.bias.data.to(x.dtype)
        return quantized_matrix_multiply(x, self.weight, self.bias)


class InitModel:

    @staticmethod
    def load_text_encoder() -> T5EncoderModel:
        print("Loading text encoder...")
        text_encoder = T5EncoderModel.from_pretrained(
            "city96/t5-v1_1-xxl-encoder-bf16",
            revision="1b9c856aadb864af93c1dcdc226c2774fa67bc86",
            torch_dtype=torch.bfloat16,
        )
        return text_encoder.to(memory_format=torch.channels_last)

    @staticmethod
    def load_transformer(trans_path: str) -> FluxTransformer2DModel:
        print("Loading transformer model...")
        transformer = FluxTransformer2DModel.from_pretrained(
            trans_path,
            torch_dtype=torch.bfloat16,
            use_safetensors=False,
        )
        return transformer.to(memory_format=torch.channels_last)



def load_pipeline() -> Pipeline:



    transformer_path = os.path.join(HF_HUB_CACHE, "models--MyApricity--Flux_Transformer_float8/snapshots/66c5f182385555a00ec90272ab711bb6d3c197db")
    transformer = InitModel.load_transformer(transformer_path)
 

    pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
                        revision=REVISION, 
                        transformer=transformer, 
                        torch_dtype=torch.bfloat16)
    pipeline.to("cuda")

    try:
        # Enable some options for better vae
        pipeline.enable_vae_slicing()
        pipeline.enable_vae_tiling()
        torch.nn.LinearLayer = QuantizedLinearLayer
    except:
        print("Debug here")

    try:
        pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)

    except:
        print("nothing")
        

    ps = [
        "overgross, mandative, inventful, braunite, penneeck",
        "melanogen, endosome, apical, polymyodous, ",
        "buffer, cutie, buttinsky, prototrophic",
        "puzzlehead",
    ]
    
    for warmprompt in ps:
        pipeline(prompt=warmprompt, 
                        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:

    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    # remove cache here for better result
    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]