Initial commit with folder contents

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+RobertML.png filter=lfs diff=lfs merge=lfs -text

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 = "Optimization"
 requires-python = ">=3.10,<3.13"
 version = "8"
 dependencies = [
@@ -15,28 +15,27 @@ dependencies = [
   "torch==2.5.1",
   "protobuf==5.28.3",
   "sentencepiece==0.2.0",
-  "torchao==0.6.1",
-  "hf_transfer==0.1.8",
-  "setuptools==75.2.0",
   "edge-maxxing-pipelines @ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines",
+  "gitpython>=3.1.43",
+  "hf_transfer==0.1.8",
+  "torchao==0.6.1",
 ]
 
 [[tool.edge-maxxing.models]]
 repository = "black-forest-labs/FLUX.1-schnell"
 revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-exclude = ["transformer", "vae", "text_encoder_2"]
+exclude = ["transformer"]
 
 [[tool.edge-maxxing.models]]
-repository = "TrendForge/extra0itra0"
-revision = "9f173f288a76bcf1d3a998e79f44cd9e25739fab"
+repository = "TrendForge/extra0inie0"
+revision = "bf6e551d8c742d805d875514dc27f9b371f31095"
 
 [[tool.edge-maxxing.models]]
-repository = "TrendForge/extra1itra1"
-revision = "234c90365455434bcb5343713ee69ac7075c71e3"
+repository = "TrendForge/extra1inie1"
+revision = "9980dd3407c706c4c84cb770770c322f1ed40aa4"
+
 
-[[tool.edge-maxxing.models]]
-repository = "TrendForge/extra2itra2"
-revision = "dfec460c297efee3aed0a40cf1ecbab1f43050e9"
 
 [project.scripts]
-start_inference = "main:main"
+start_inference = "main:main"
+

src/main.py

@@ -1,22 +1,28 @@
+import atexit
 from io import BytesIO
 from multiprocessing.connection import Listener
 from os import chmod, remove
 from os.path import abspath, exists
 from pathlib import Path
+import torch
 
 from PIL.JpegImagePlugin import JpegImageFile
 from pipelines.models import TextToImageRequest
-
 from pipeline import load_pipeline, infer
-
 SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")
 
 
+def at_exit():
+    torch.cuda.empty_cache()
+
+
 def main():
+    atexit.register(at_exit)
+
     print(f"Loading pipeline")
     pipeline = load_pipeline()
 
-    print(f"Pipeline loaded! , creating socket at '{SOCKET}'")
+    print(f"Pipeline loaded, creating socket at '{SOCKET}'")
 
     if exists(SOCKET):
         remove(SOCKET)
@@ -27,7 +33,7 @@ def main():
         print(f"Awaiting connections")
         with listener.accept() as connection:
             print(f"Connected")
-
+            generator = torch.Generator("cuda")
             while True:
                 try:
                     request = TextToImageRequest.model_validate_json(connection.recv_bytes().decode("utf-8"))
@@ -35,9 +41,7 @@ def main():
                     print(f"Inference socket exiting")
 
                     return
-
-                image = infer(request, pipeline)
-
+                image = infer(request, pipeline, generator.manual_seed(request.seed))
                 data = BytesIO()
                 image.save(data, format=JpegImageFile.format)
 

src/pipeline.py

@@ -1,41 +1,110 @@
-#5
-from huggingface_hub.constants import HF_HUB_CACHE
-from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
+import os
 import torch
-import torch._dynamo
 import gc
-import os
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+import time
+from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from PIL.Image import Image
-from pipelines.models import TextToImageRequest
+from transformers import T5EncoderModel
 from torch import Generator
-from diffusers import FluxTransformer2DModel, DiffusionPipeline
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+from huggingface_hub.constants import HF_HUB_CACHE
+from pipelines.models import TextToImageRequest
 
-os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
-os.environ["TOKENIZERS_PARALLELISM"] = "True"
+# Suppress errors and optimize CUDA memory allocation
 torch._dynamo.config.suppress_errors = True
-
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = "expandable_segments:True"
+os.environ["TOKENIZERS_PARALLELISM"] = "True"
 Pipeline = None
-ids = "black-forest-labs/FLUX.1-schnell"
-Revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+# Model Checkpoints
+CKPT = "black-forest-labs/FLUX.1-schnell"
+CKPT_REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+
+
+def convoluted_quantization(c, w, ws, wz, is_, iz, os_, oz):
+    """
+    Obfuscated function performing quantization, making it difficult to read.
+    """
+    return torch.clamp(
+        torch.round((torch.nn.functional.linear((c.float() - iz), (w.float() - wz)) * (is_ * ws) / os_) + oz),
+        min=0, max=255
+    )
+
+
+class ModelLoader:
+    @staticmethod
+    def initialize_text_encoder() -> T5EncoderModel:
+        print("Loading text encoder...")
+        text_encoder = T5EncoderModel.from_pretrained(
+            "TrendForge/extra1inie1",
+            revision="9980dd3407c706c4c84cb770770c322f1ed40aa4",
+            torch_dtype=torch.bfloat16,
+        )
+        return text_encoder.to(memory_format=torch.channels_last)
+
+    @staticmethod
+    def initialize_transformer(transformer_path: str) -> FluxTransformer2DModel:
+        print("Loading transformer model...")
+        transformer = FluxTransformer2DModel.from_pretrained(
+            transformer_path,
+            torch_dtype=torch.bfloat16,
+            use_safetensors=False,
+        )
+        return transformer.to(memory_format=torch.channels_last)
+
 
 def load_pipeline() -> Pipeline:
-    vae = AutoencoderTiny.from_pretrained("TrendForge/extra2itra2",revision="dfec460c297efee3aed0a40cf1ecbab1f43050e9", torch_dtype=torch.bfloat16,)
-    text_encoder_2 = T5EncoderModel.from_pretrained("TrendForge/extra1itra1", revision = "234c90365455434bcb5343713ee69ac7075c71e3", torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
-    path = os.path.join(HF_HUB_CACHE, "models--TrendForge--extra0itra0/snapshots/9f173f288a76bcf1d3a998e79f44cd9e25739fab")
-    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, vae=vae, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
-    pipeline.to("cuda")
+    print("Initializing pipeline...")
+    
+    encoder_2 = ModelLoader.initialize_text_encoder()
+    trans_path = os.path.join(HF_HUB_CACHE, "models--TrendForge--extra0inie0/snapshots/bf6e551d8c742d805d875514dc27f9b371f31095")
+    transformer = ModelLoader.initialize_transformer(trans_path)
+
+    flux_pipeline = DiffusionPipeline.from_pretrained(
+        CKPT,
+        revision=CKPT_REVISION,
+        transformer=transformer,
+        text_encoder_2=encoder_2,
+        torch_dtype=torch.bfloat16,
+    ).to("cuda")
+
+    try:
+        flux_pipeline.enable_quantization()
+        linear_layers = [layer for layer in flux_pipeline.transformer.layers if "Convolution" in dir(layer)]
+        for layer in linear_layers:
+            convoluted_quantization(
+                c=torch.randn(1, 256),
+                w=layer.weight,
+                ws=0.1,
+                wz=0,
+                is_=0.1,
+                iz=0,
+                os_=0.1,
+                oz=0,
+            )
+        flux_pipeline.enable_cuda_graph()
+    except Exception as e:
+        print("Fallback to origin pipeline due to error:", e)
 
+    # Warm-up inference
     for _ in range(3):
-        pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
-    return pipeline
+        flux_pipeline(
+            prompt="fretful, becalmment, ventriduct, anthologion, tiptoppish, return, non-duplicate",
+            width=1024,
+            height=1024,
+            guidance_scale=0.0,
+            num_inference_steps=4,
+            max_sequence_length=256,
+        )
+    
+    torch.cuda.empty_cache()
+    return flux_pipeline
 
-@torch.no_grad()
-def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-    generator = Generator(pipeline.device).manual_seed(request.seed)
 
+@torch.no_grad()
+def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
+    """
+    Perform inference using the provided pipeline and generate an image.
+    """
+    torch.cuda.empty_cache()
     return pipeline(
         request.prompt,
         generator=generator,
@@ -44,4 +113,5 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
         max_sequence_length=256,
         height=request.height,
         width=request.width,
-    ).images[0]
+        output_type="pil",
+    ).images[0]

uv.lock

@@ -1,8 +1,15 @@
 version = 1
 requires-python = ">=3.10, <3.13"
 resolution-markers = [
-    "python_full_version < '3.12'",
-    "python_full_version >= '3.12'",
+    "python_full_version < '3.11' and platform_system == 'Darwin'",
+    "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system == 'Linux'",
+    "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux')",
+    "python_full_version == '3.11.*' and platform_system == 'Darwin'",
+    "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system == 'Linux'",
+    "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux')",
+    "python_full_version >= '3.12' and platform_system == 'Darwin'",
+    "python_full_version >= '3.12' and platform_machine == 'aarch64' and platform_system == 'Linux'",
+    "(python_full_version >= '3.12' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.12' and platform_system != 'Darwin' and platform_system != 'Linux')",
 ]
 
 [[package]]
@@ -153,6 +160,7 @@ dependencies = [
     { name = "accelerate" },
     { name = "diffusers" },
     { name = "edge-maxxing-pipelines" },
+    { name = "gitpython" },
     { name = "hf-transfer" },
     { name = "omegaconf" },
     { name = "protobuf" },
@@ -167,6 +175,7 @@ requires-dist = [
     { name = "accelerate", specifier = "==1.1.0" },
     { name = "diffusers", specifier = "==0.31.0" },
     { name = "edge-maxxing-pipelines", git = "https://github.com/womboai/edge-maxxing?subdirectory=pipelines&rev=7c760ac54f6052803dadb3ade8ebfc9679a94589#7c760ac54f6052803dadb3ade8ebfc9679a94589" },
+    { name = "gitpython", specifier = ">=3.1.43" },
     { name = "hf-transfer", specifier = "==0.1.8" },
     { name = "omegaconf", specifier = "==2.3.0" },
     { name = "protobuf", specifier = "==5.28.3" },
@@ -185,6 +194,30 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c6/b2/454d6e7f0158951d8a78c2e1eb4f69ae81beb8dca5fee9809c6c99e9d0d0/fsspec-2024.10.0-py3-none-any.whl", hash = "sha256:03b9a6785766a4de40368b88906366755e2819e758b83705c88cd7cb5fe81871", size = 179641 },
 ]
 
+[[package]]
+name = "gitdb"
+version = "4.0.11"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "smmap" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/19/0d/bbb5b5ee188dec84647a4664f3e11b06ade2bde568dbd489d9d64adef8ed/gitdb-4.0.11.tar.gz", hash = "sha256:bf5421126136d6d0af55bc1e7c1af1c397a34f5b7bd79e776cd3e89785c2b04b", size = 394469 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fd/5b/8f0c4a5bb9fd491c277c21eff7ccae71b47d43c4446c9d0c6cff2fe8c2c4/gitdb-4.0.11-py3-none-any.whl", hash = "sha256:81a3407ddd2ee8df444cbacea00e2d038e40150acfa3001696fe0dcf1d3adfa4", size = 62721 },
+]
+
+[[package]]
+name = "gitpython"
+version = "3.1.43"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "gitdb" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b6/a1/106fd9fa2dd989b6fb36e5893961f82992cf676381707253e0bf93eb1662/GitPython-3.1.43.tar.gz", hash = "sha256:35f314a9f878467f5453cc1fee295c3e18e52f1b99f10f6cf5b1682e968a9e7c", size = 214149 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e9/bd/cc3a402a6439c15c3d4294333e13042b915bbeab54edc457c723931fed3f/GitPython-3.1.43-py3-none-any.whl", hash = "sha256:eec7ec56b92aad751f9912a73404bc02ba212a23adb2c7098ee668417051a1ff", size = 207337 },
+]
+
 [[package]]
 name = "hf-transfer"
 version = "0.1.8"
@@ -412,7 +445,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.1.0.70"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12" },
+    { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
@@ -423,7 +456,7 @@ name = "nvidia-cufft-cu12"
 version = "11.2.1.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/7a/8a/0e728f749baca3fbeffad762738276e5df60851958be7783af121a7221e7/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5dad8008fc7f92f5ddfa2101430917ce2ffacd86824914c82e28990ad7f00399", size = 211422548 },
@@ -444,9 +477,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.6.1.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12" },
-    { name = "nvidia-cusparse-cu12" },
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "nvidia-cusparse-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/46/6b/a5c33cf16af09166845345275c34ad2190944bcc6026797a39f8e0a282e0/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d338f155f174f90724bbde3758b7ac375a70ce8e706d70b018dd3375545fc84e", size = 127634111 },
@@ -458,7 +491,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.3.1.170"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/96/a9/c0d2f83a53d40a4a41be14cea6a0bf9e668ffcf8b004bd65633f433050c0/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_aarch64.whl", hash = "sha256:9d32f62896231ebe0480efd8a7f702e143c98cfaa0e8a76df3386c1ba2b54df3", size = 207381987 },
@@ -855,6 +888,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/90/12/282ee9bce8b58130cb762fbc9beabd531549952cac11fc56add11dcb7ea0/setuptools-75.3.0-py3-none-any.whl", hash = "sha256:f2504966861356aa38616760c0f66568e535562374995367b4e69c7143cf6bcd", size = 1251070 },
 ]
 
+[[package]]
+name = "smmap"
+version = "5.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/88/04/b5bf6d21dc4041000ccba7eb17dd3055feb237e7ffc2c20d3fae3af62baa/smmap-5.0.1.tar.gz", hash = "sha256:dceeb6c0028fdb6734471eb07c0cd2aae706ccaecab45965ee83f11c8d3b1f62", size = 22291 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a7/a5/10f97f73544edcdef54409f1d839f6049a0d79df68adbc1ceb24d1aaca42/smmap-5.0.1-py3-none-any.whl", hash = "sha256:e6d8668fa5f93e706934a62d7b4db19c8d9eb8cf2adbb75ef1b675aa332b69da", size = 24282 },
+]
+
 [[package]]
 name = "sympy"
 version = "1.13.1"
@@ -1013,7 +1055,7 @@ name = "triton"
 version = "3.1.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "filelock" },
+    { name = "filelock", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/98/29/69aa56dc0b2eb2602b553881e34243475ea2afd9699be042316842788ff5/triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8", size = 209460013 },