from diffusers import FluxPipeline, AutoencoderKL, FluxTransformer2DModel #AutoencoderTiny from diffusers.image_processor import VaeImageProcessor from diffusers.schedulers import FlowMatchEulerDiscreteScheduler import torch.nn.functional as F from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel import torch import torch._dynamo import gc from PIL import Image as img from PIL.Image import Image from pipelines.models import TextToImageRequest from torch import Generator import time from diffusers import FluxTransformer2DModel, DiffusionPipeline import torch.nn as nn # from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight #PerRow, import os os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01" os.environ["HUGGINGFACE_HUB_TOKEN"] = "" Pipeline = None # def w8_a16_forward(weight, input, scales, bias=None): # casted_weights = weight.to(input.dtype) # output = F.linear(input, casted_weights) * scales # overhead # if bias is not None: # output = output + bias # return output # class W8A16LinearLayer(nn.Module): # def __init__(self, in_features, out_features, bias=True, dtype=torch.float32): # super().__init__() # self.register_buffer( # "int8_weights", # torch.randint(-128, 127, (out_features, in_features), dtype=torch.int8)) # self.register_buffer("scales", torch.randn((out_features), dtype=dtype)) # if bias: # self.register_buffer("bias", torch.randn((1, out_features), dtype=dtype)) # def quantize(self, weights): # w_fp32 = weights.clone().to(torch.float32) # scales = w_fp32.abs().max(dim=-1).values / 127 # scales = scales.to(weights.dtype) # int8_weights = torch.round(weights/scales.unsqueeze(1)).to(torch.int8) # self.int8_weights = int8_weights # self.scales = scales # self.bias = None # def forward(self, input): # return w8_a16_forward(self.int8_weights, input, self.scales, self.bias) class W8A16LinearLayer(nn.Module): def __init__(self, in_features, out_features, bias=True, dtype=torch.float32): super().__init__() self.weight = nn.Parameter(torch.randn(out_features, in_features, dtype=dtype)) self.weight.requires_grad = False if bias: self.bias = nn.Parameter(torch.randn(1, out_features, dtype=dtype)) self.scales = nn.Parameter(torch.randn(out_features, dtype=dtype)) def quantize(self, weights): w_fp32 = weights.clone().to(torch.float32) scales = w_fp32.abs().max(dim=-1).values / 127 scales = scales.to(weights.dtype) self.weight.data = torch.round(weights/scales.unsqueeze(1)).to(torch.int8) self.scales.data = scales def forward(self, input): casted_weights = self.weight.to(input.dtype) output = F.linear(input, casted_weights) * self.scales if self.bias is not None: output = output + self.bias return output # def replace_linear_with_target_and_quantize(module, target_class, module_name_to_exclude): # # with open("/root/.cache/huggingface/hub/output_layers.txt", "a") as f: # for name, child in module.named_children(): # if isinstance(child, nn.Linear) and ( 'add_k_proj' in name or 'add_v_proj' in name or 'add_q_proj' in name ): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or # old_bias = child.bias # old_weight = child.weight # new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype) # new_module.quantize(old_weight) # delattr(module, name) # setattr(module, name, new_module) # if old_bias is not None: # getattr(module, name).bias = old_bias # # # Print the replaced layer name and calculate the change in size # # old_size = old_weight.numel() * old_weight.element_size() # # new_size = new_module.int8_weights.numel() * new_module.int8_weights.element_size() # # f.write(f"Replaced layer: {name}" + f" Size reduction: {old_size} bytes -> {new_size} bytes ({(old_size - new_size) / old_size * 100:.2f}% reduction)") # else: # # Recursively call the function for nested modules # replace_linear_with_target_and_quantize(child, target_class, module_name_to_exclude) def replace_linear_with_target_and_quantize(module, target_class, module_name_to_exclude): # with open("/root/.cache/huggingface/hub/output_layers.txt", "a") as f: for name in list(module._modules.keys()): child = module._modules[name] if isinstance(child, nn.Linear) and ( 'add_k_proj' in name or 'add_v_proj' in name or 'add_q_proj' in name ): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or old_bias = child.bias old_weight = child.weight new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype) new_module.quantize(old_weight) delattr(module, name) setattr(module, name, new_module) if old_bias is not None: getattr(module, name).bias = old_bias # # Print the replaced layer name and calculate the change in size # old_size = old_weight.numel() * old_weight.element_size() # new_size = new_module.int8_weights.numel() * new_module.int8_weights.element_size() # f.write(f"Replaced layer: {name}" + f" Size reduction: {old_size} bytes -> {new_size} bytes ({(old_size - new_size) / old_size * 100:.2f}% reduction)") else: # Recursively call the function for nested modules replace_linear_with_target_and_quantize(child, target_class, module_name_to_exclude) ckpt_id = "black-forest-labs/FLUX.1-schnell" def empty_cache(): start = time.time() gc.collect() torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() def load_pipeline() -> Pipeline: empty_cache() dtype, device = torch.bfloat16, "cuda" text_encoder_2 = T5EncoderModel.from_pretrained( "city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16 ) vae=AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=dtype) # vae = torch.load('/root/.cache/huggingface/hub/compiled_vae.pth') # transformer = FluxTransformer2DModel.from_pretrined("manbeast3b/transfomer-flux-schnell-int8") # torch_dtype=dtype pipeline = DiffusionPipeline.from_pretrained( ckpt_id, vae=vae, text_encoder_2 = text_encoder_2, # transformer=transformer, torch_dtype=dtype, ) # quantize_(pipeline.transformer, float8_dynamic_activation_float8_weight()) torch.backends.cudnn.benchmark = True torch.backends.cudnn.deterministic = False # torch.set_deterministic_debug_mode(0) torch.backends.cuda.matmul.allow_tf32 = True # torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True torch.cuda.set_memory_growth(True) torch.cuda.set_per_process_memory_fraction(0.99) pipeline.text_encoder.to(memory_format=torch.channels_last) pipeline.transformer.to(memory_format=torch.channels_last) # replace_linear_with_target_and_quantize(pipeline.transformer, W8A16LinearLayer, []) # pipeline.transformer.save_pretrained("manbeast3b/transfomer-flux-schnell-int8-new", push_to_hub=True, token="") # pipeline.transformer.save_pretrained("/root/.cache/huggingface/hub/transformer-flux") # exit() pipeline.vae.to(memory_format=torch.channels_last) pipeline.vae = torch.compile(pipeline.vae) # torch.save(pipeline.vae, '/root/.cache/huggingface/hub/compiled_vae.pth') # exit() pipeline._exclude_from_cpu_offload = ["vae"] pipeline.enable_sequential_cpu_offload() for _ in range(2): pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256) return pipeline @torch.inference_mode() def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image: torch.cuda.reset_peak_memory_stats() generator = Generator("cuda").manual_seed(request.seed) image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0] return(image)