i

pyproject.toml

@@ -23,20 +23,14 @@ dependencies = [
 ]
 
 [[tool.edge-maxxing.models]]
-repository = "black-forest-labs/FLUX.1-schnell"
-revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+repository = "MyApricity/FLUX_OPT_SCHNELL_1.2"
+revision = "488528b6f815bff1bbc747cf1e0947c77c544665"
 
 [[tool.edge-maxxing.models]]
 repository = "city96/t5-v1_1-xxl-encoder-bf16"
 revision = "1b9c856aadb864af93c1dcdc226c2774fa67bc86"
 
-[[tool.edge-maxxing.models]]
-repository = "MyApricity/Vae_Only"
-revision = "a47d57702caf8ff0c0e21d30b93f9d3297b81920"
 
-[[tool.edge-maxxing.models]]
-repository = "MyApricity/Flux_Transformer_float8"
-revision = "66c5f182385555a00ec90272ab711bb6d3c197db"
 
 [project.scripts]
 start_inference = "main:main"

src/pipeline.py

@@ -8,7 +8,8 @@ import transformers
 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
+# ApricityApricityApricityApricityApricityApricityApricityApricityApricityApricityApricity
+
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 
@@ -19,60 +20,49 @@ from optimum.quanto import requantize
 import json
 
 
+# ApricityApricityApricityApricityApricityApricityApricityApricityApricityApricityApricity
 
 
 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"
+ckpt_root = "MyApricity/FLUX_OPT_SCHNELL_1.2"
+revision_root = "488528b6f815bff1bbc747cf1e0947c77c544665"
 Pipeline = None
-
+use_com = False
 
 import torch
 import math
 from typing import Dict, Any
 
 def remove_cache():
-    gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
+    gc.collect()
     torch.cuda.reset_peak_memory_stats()
 
 
-class InitializingModel:
 
-    @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)
+def text_t5_loader() -> 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)
 
-
-class CompileTransformerDiffusion:
+class StableDiffusionTransformerCompile:
     def __init__(self, pipeline, optimize=False):
         self.pipeline = pipeline
         self.optimize = optimize
         if self.optimize:
-            self._compile_model()
+            self.model_compiling()
     
-    def _compile_model(self):
-        print("Compiling transformer model for optimized diffusion...")
+    def model_compiling(self):
+        # Staff doing here
         self.pipeline.unet = torch.compile(self.pipeline.unet)
     
     def __call__(self, *args, **kwargs):
@@ -80,39 +70,55 @@ class CompileTransformerDiffusion:
     
 def load_pipeline() -> Pipeline:
 
+    text_t5_encoder = text_t5_loader()
 
-    base_transformer_path = os.path.join(HF_HUB_CACHE, "models--MyApricity--Flux_Transformer_float8/snapshots/66c5f182385555a00ec90272ab711bb6d3c197db")
-    base_transformer = InitializingModel.load_transformer(base_transformer_path)
+    transformer_path__ = os.path.join(HF_HUB_CACHE, "models--MyApricity--FLUX_OPT_SCHNELL_1.2/snapshots/488528b6f815bff1bbc747cf1e0947c77c544665")
     
-    text_encoder_2 = InitializingModel.load_text_encoder()
+    transformer__ = FluxTransformer2DModel.from_pretrained(transformer_path__, torch_dtype=torch.bfloat16, use_safetensors=False)
 
+    try:
+        pipeline = DiffusionPipeline.from_pretrained(ckpt_root, 
+                            revision=revision_root, 
+                            transformer=transformer__, 
+                            torch_dtype=torch.bfloat16)
+
+    except:
+        pipeline = DiffusionPipeline.from_pretrained(ckpt_root, 
+                            revision=revision_root,
+                            torch_dtype=torch.bfloat16)
 
-    pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
-                        revision=REVISION, 
-                        transformer=base_transformer, 
-                        text_encoder_2=text_encoder_2, 
-                        torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
+
     try:
-        pipeline.disable_vae_slice()
-        compiled_pipeline = CompileTransformerDiffusion(pipeline, optimize=False)
+        compiled_pipeline = StableDiffusionTransformerCompile(pipeline, optimize=False)
+
+        if use_com:
+            pipeline = compiled_pipeline
+        else:
+            print("Currently not compling affectively")
+
+        pipeline.disable_vae_compress()
+        pipeline.text_encoder_2 = text_t5_encoder
+
     except:
-        print("Stay safe here pipeline")
+        print("pipeline")
         
 
-    promts_listing = [
-        "sellate, Tremellales, thro, albescent",
-        "must return non duplicate",
-        "albaspidin, pillmonger, palaeocrystalline"
-    ]
-    
-    for p in promts_listing:
-        pipeline(prompt=p, 
-                        width=1024,
-                        height=1024,
-                        guidance_scale=0.0,
-                        num_inference_steps=4,
-                        max_sequence_length=256)
+    prompt_1 = "albaspidin, pillmonger, palaeocrystalline"
+    pipeline(prompt=prompt_1, 
+                    width=1024,
+                    height=1024,
+                    guidance_scale=0.0,
+                    num_inference_steps=4,
+                    max_sequence_length=256)
+
+    prompt_2 = "obe, kilometrage, circuition"
+    pipeline(prompt=prompt_2, 
+                    width=1024,
+                    height=1024,
+                    guidance_scale=0.0,
+                    num_inference_steps=4,
+                    max_sequence_length=256)
 
     return pipeline