u

src/pipeline.py

@@ -1,14 +1,10 @@
-# FLux Optimization Pipeline
+# asfsdgdvsdgtwtgfgfsgsgdsvxcvsgsg
 import os
 import torch
 import torch._dynamo
 import gc
 
-
 from huggingface_hub.constants import HF_HUB_CACHE
-from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
-
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 
@@ -18,7 +14,7 @@ from pipelines.models import TextToImageRequest
 from optimum.quanto import requantize
 import json
 import transformers
-
+# ITs meeeeeeeeeeeeeeeeeeeeeeeeeeeeee
 
 
 torch._dynamo.config.suppress_errors = True
@@ -47,6 +43,20 @@ def load_quanto_text_encoder_2(text_repo_path):
     requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cuda"))
     return text_encoder_2
 
+class TransformerDiffusionCompiler:
+    def __init__(self, diffusion_pipeline, activate_optimization=False):
+        self.pipeline = diffusion_pipeline
+        self.optimization = activate_optimization
+        if self.optimization:
+            self._enhance_model()
+    
+    def _enhance_model(self):
+        print("Applying advanced optimizations to the transformer pipeline...")
+        self.pipeline.unet = torch.compile(self.pipeline.unet, mode="reduce-overhead")
+    
+    def execute(self, *params, **kwargs):
+        return self.pipeline(*params, **kwargs)
+
 def load_pipeline() -> Pipeline:
 
     origin_vae = AutoencoderTiny.from_pretrained("RichardWilliam/XULF_Vae",
@@ -70,11 +80,12 @@ def load_pipeline() -> Pipeline:
         # pipeline.text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None)
         pipeline.enable_cuda_graph(type="max-autotune")
         pipeline.text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None)
+        pipeline = TransformerDiffusionCompiler(pipeline, activate_optimization=True)
     except:
         print("Something wrong here")
 
     for __ in range(3):
-        pipeline(prompt="schoenobatist, halisteresis, chronometric, hallucinative", 
+        pipeline(prompt="whensoever, lodger, Locarnize, hippology, harakeke", 
                         width=1024, 
                         height=1024, 
                         guidance_scale=0.0,