# !pip install "git+https://github.com/ishan-modi/diffusers.git@add-trtquant-backend#egg=diffusers[nvidia_modelopt]"
import torch
from tqdm import tqdm
from diffusers import SanaPipeline, SanaTransformer2DModel, StableDiffusion3Pipeline, , SD3Transformer2DModel
from diffusers.quantizers.quantization_config import NVIDIAModelOptConfig
checkpoint = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
model_cls = SanaTransformer2DModel
pipe_cls = SanaPipeline
# checkpoint = "stabilityai/stable-diffusion-3-medium-diffusers"
# model_cls = SD3Transformer2DModel
# pipe_cls = StableDiffusion3Pipeline
input = {"prompt":"A capybara holding a sign that reads Hello World", "num_inference_steps":28, "guidance_scale":3.5}
quant_config_fp8 = {"quant_type": "FP8", "quant_method": "modelopt"}
quant_config_int8 = {"quant_type": "INT8", "quant_method": "modelopt"}
quant_config_int4 = {"quant_type": "INT4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "modules_to_not_convert": ["conv", "patch_embed"]}
quant_config_nf4 = {"quant_type": "NF4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "scale_block_quantize": 8, "scale_channel_quantize": -1, "modules_to_not_convert": ["conv"]}
quant_config_nvfp4 = {"quant_type": "NVFP4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "modules_to_not_convert": ["conv"]}
def test_quantization(config, checkpoint, model_cls):
 quant_config = NVIDIAModelOptConfig(**config)
 print(quant_config.get_config_from_quant_type())
 quant_model = model_cls.from_pretrained(checkpoint, subfolder="transformer", quantization_config=quant_config, torch_dtype=torch.bfloat16).to('cuda')
 return quant_model
def test_quant_inference(model, input, pipe_cls, iter=1):
 inference_memory = 0
 for _ in tqdm(range(iter)):
 with torch.no_grad():
 output = pipe_cls.from_pretrained(checkpoint, transformer=model, torch_dtype=torch.bfloat16).to('cuda')(**input).images[0]
 inference_memory += torch.cuda.max_memory_allocated()
 inference_memory /= iter
 output.save("test.png")
 print("Inference Memory: ", inference_memory / 1e6)
 torch.cuda.empty_cache()
 torch.cuda.reset_max_memory_allocated()
test_quant_inference(test_quantization(quant_config_fp8, checkpoint, model_cls), input, pipe_cls)
# test_quant_inference(test_quantization(quant_config_int8, checkpoint, model_cls), input, pipe_cls)
# test_quant_inference(test_quantization(quant_config_int4, checkpoint, model_cls), input, pipe_cls)
# test_quant_inference(test_quantization(quant_config_nf4, checkpoint, model_cls), input, pipe_cls)
# test_quant_inference(test_quantization(quant_config_nvfp4, checkpoint, model_cls), input, pipe_cls)
# test_quant_inference(model_cls.from_pretrained(checkpoint, subfolder="transformer", torch_dtype=torch.bfloat16).to('cuda'), input, pipe_cls)