u

pyproject.toml

@@ -23,12 +23,9 @@ dependencies = [
 ]
 
 [[tool.edge-maxxing.models]]
-repository = "black-forest-labs/FLUX.1-schnell"
-revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-
-[[tool.edge-maxxing.models]]
-repository = "RichardWilliam/XULF_T5_bf16"
-revision = "63a3d9ef7b586655600ac9bd4e4747d038237761"
+repository = "RichardWilliam/FullyFLUXSCH"
+revision = "c5f4f70c6cb9228a9c258799aadc660dde417af6"
+exclude = ["transformer"]
 
 [[tool.edge-maxxing.models]]
 repository = "RichardWilliam/XULF_Vae"

src/pipeline.py

@@ -1,9 +1,14 @@
-# asfsdgdvsdgtwtgfgfsgsgdsvxcvsgsg
+# eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee
 import os
 import torch
 import torch._dynamo
 import gc
 
+torch._dynamo.config.suppress_errors = True
+os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
+os.environ["TOKENIZERS_PARALLELISM"] = "True"
+
+
 from huggingface_hub.constants import HF_HUB_CACHE
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
@@ -14,24 +19,36 @@ from pipelines.models import TextToImageRequest
 from optimum.quanto import requantize
 import json
 import transformers
-# ITs meeeeeeeeeeeeeeeeeeeeeeeeeeeeee
+from functools import wraps
+
 
 
 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"
+MAIN_ID = "RichardWilliam/FullyFLUXSCH"
+REV = "c5f4f70c6cb9228a9c258799aadc660dde417af6"
 Pipeline = None
 apply_quanto=1
 
-def reset_cache():
+def to_hell():
     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()
 
+def error_handler(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            print(f"Error in {func.__name__}: {str(e)}")
+            return None
+    return wrapper
+
+@error_handler
 def load_quanto_text_encoder_2(text_repo_path):
     with open("quantization_map.json", "r") as f:
         quantization_map = json.load(f)
@@ -43,49 +60,32 @@ 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",
-                    revision="3ee225c539465c27adadec45c6e8af50a7397b7d", 
-                    torch_dtype=torch.bfloat16)
-
-
     main_path = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_Transfomer/snapshots/6860c51af40329808f270e159a0d018559a1204f")
     origin_trans = FluxTransformer2DModel.from_pretrained(main_path,
                         torch_dtype=torch.bfloat16, 
                         use_safetensors=False).to(memory_format=torch.channels_last)
+    
     transformer = origin_trans
 
-    pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
-                        revision=REVISION,
+    pipeline = DiffusionPipeline.from_pretrained(MAIN_ID, 
+                        revision=REV,
                         transformer=transformer, 
                         torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
 
-    try:
-        # 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")
+
+    text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None)
+
+    if text_encoder_v2==None:
+        print("Something wrong")
+    else:
+        pipeline.text_encoder_2 = text_encoder_v2
 
     for __ in range(3):
-        pipeline(prompt="whensoever, lodger, Locarnize, hippology, harakeke", 
+        pipeline(prompt="I am the worst", 
                         width=1024, 
                         height=1024, 
                         guidance_scale=0.0, 
@@ -96,8 +96,7 @@ def load_pipeline() -> Pipeline:
 @torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
 
-    reset_cache()
-    torch.cuda.empty_cache()
+    to_hell()
 
     generator = Generator(pipeline.device).manual_seed(request.seed)