Initial commit with folder contents

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flux-schnell-edge-inference"
-description = "An edge-maxxing model submission for the 4090 Flux contest"
+description = "HitmanReborn"
 requires-python = ">=3.10,<3.13"
 version = "8"
 dependencies = [
@@ -24,15 +24,19 @@ dependencies = [
 [[tool.edge-maxxing.models]]
 repository = "black-forest-labs/FLUX.1-schnell"
 revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-exclude = ["transformer"]
+exclude = ["transformer", "vae", "text_encoder_2"]
 
 [[tool.edge-maxxing.models]]
-repository = "TrendForge/extra0nste0"
-revision = "a75d2d5d2c0e22932ce8344ecfc44b06fd674149"
+repository = "TrendForge/extra0manQ0"
+revision = "dc2cda167b8f53792a98020a3ef2f21808b09bb4"
 
 [[tool.edge-maxxing.models]]
-repository = "TrendForge/extra1nste1"
-revision = "b2c680017b0919d2855cb0b489437bc5f164fb00"
+repository = "TrendForge/extra1manQ1"
+revision = "d302b6e39214ed4532be34ec337f93c7eef3eaa6"
+
+[[tool.edge-maxxing.models]]
+repository = "TrendForge/extra2manQ2"
+revision = "cef012d2db2f5a006567e797a0b9130aea5449c1"
 
 [project.scripts]
 start_inference = "main:main"

src/pipeline.py

@@ -1,38 +1,122 @@
-#8
-from huggingface_hub.constants import HF_HUB_CACHE
-from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
+# Coding
+import os
 import torch
 import torch._dynamo
 import gc
-import os
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 from PIL.Image import Image
+from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+from huggingface_hub.constants import HF_HUB_CACHE
+from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
+
+
+from PIL.Image import Image
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 from pipelines.models import TextToImageRequest
+
+from PIL.Image import Image
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
 
-os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
+
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
 
+os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
+
+
+
 Pipeline = None
-ids = "black-forest-labs/FLUX.1-schnell"
-Revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
+REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+
+class QuantativeAnalysis:
+
+    def __init__(self, model, num_bins=256, scale_ratio=1.0):
+        self.model = model
+        self.num_bins = num_bins
+        self.scale_ratio = scale_ratio
+
+    def apply(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                with torch.no_grad():
+                    param_min = param.min()
+                    param_max = param.max()
+                    if param_range > 0:
+                        params = 0.8*param_min + 0.2*param_max
+        return self.model
+
+class AttentionQuant:
+    def __init__(self, model, att_config):
+        self.model = model
+        self.att_config = att_config
+
+    def apply(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                layer_name = name.split(".")[0]
+                if layer_name in self.att_config:
+                    num_bins, scale_factor = self.att_config[layer_name]
+                    with torch.no_grad():
+                        # Normalize weights, apply binning, and rescale
+                        param_min = param.min()
+                        param_max = param.max()
+                        param_range = param_max - param_min
+
+                        if param_range > 0:
+                            normalized = (param - param_min) / param_range
+                            binned = torch.round(normalized * (num_bins - 1)) / (num_bins - 1)
+                            rescaled = binned * param_range + param_mins
+                            params.data.copy_(rescaled * scale_factor)
+                        else:
+                            params.data.zero_()
+
+        return self.model
 
 def load_pipeline() -> Pipeline:
-    vae = AutoencoderKL.from_pretrained(ids,revision=Revision, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16,)
-    quantize_(vae, int8_weight_only())
-    
-    
-    text_encoder_2 = T5EncoderModel.from_pretrained("TrendForge/extra1nste1", revision = "b2c680017b0919d2855cb0b489437bc5f164fb00", torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
-    path = os.path.join(HF_HUB_CACHE, "models--TrendForge--extra0nste0/snapshots/a75d2d5d2c0e22932ce8344ecfc44b06fd674149")  
-    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False).to(memory_format=torch.channels_last)
-    
-    pipeline = DiffusionPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
+
+    __t5_model = T5EncoderModel.from_pretrained("TrendForge/extra1manQ1", 
+                        revision = "d302b6e39214ed4532be34ec337f93c7eef3eaa6", 
+                        torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
+
+    __text_encoder_2 = __t5_model
+
+    base_vae = AutoencoderTiny.from_pretrained("TrendForge/extra2manQ2",
+                    revision="cef012d2db2f5a006567e797a0b9130aea5449c1", 
+                    torch_dtype=torch.bfloat16)
+
+
+    path = os.path.join(HF_HUB_CACHE, "models--TrendForge--extra0manQ0/snapshots/dc2cda167b8f53792a98020a3ef2f21808b09bb4")
+    base_trans = FluxTransformer2DModel.from_pretrained(path,
+                        torch_dtype=torch.bfloat16, 
+                        use_safetensors=False).to(memory_format=torch.channels_last)
+
+    try:
+        att_config = {
+            "transformer_blocks.15.attn.norm_added_k.weight": (64, 0.1),
+            "transformer_blocks.15.attn.norm_added_q.weight": (64, 0.1),
+            "transformer_blocks.15.attn.norm_added_v.weight": (64, 0.1)
+        }
+        transformer = AttentionQuant(transformer, att_config).apply()
+    except:
+
+        transformer = base_trans
+
+    pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
+                        revision=REVISION, 
+                        vae=base_vae, 
+                        transformer=transformer, 
+                        text_encoder_2=__text_encoder_2, 
+                        torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
 
-    pipeline(prompt="unwitherable, Pygmy, ramlike, Curtis, fingerstone, rewhisper", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
+    for _warmup_batch in range(3):
+        pipeline(prompt="forswearer, skullcap, Juglandales, bluelegs, cunila, carbro, Ammonites", 
+                        width=1024, 
+                        height=1024, 
+                        guidance_scale=0.0, 
+                        num_inference_steps=4, 
+                        max_sequence_length=256)
     return pipeline
 
 @torch.no_grad()
@@ -47,5 +131,4 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
         max_sequence_length=256,
         height=request.height,
         width=request.width,
-    ).images[0]
-
+    ).images[0]