Initial commit with folder contents

.gitattributes

@@ -33,4 +33,3 @@ 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 by RobertML for the 4090 Flux contest"
+description = "An edge-maxxing model submission for the 4090 Flux contest"
 requires-python = ">=3.10,<3.13"
 version = "8"
 dependencies = [
@@ -15,30 +15,15 @@ dependencies = [
   "torch==2.5.1",
   "protobuf==5.28.3",
   "sentencepiece==0.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",
+  "hf_transfer==0.1.8",
+  "setuptools==75.2.0",
+  "edge-maxxing-pipelines @ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines",
 ]
 
 [[tool.edge-maxxing.models]]
-repository = "black-forest-labs/FLUX.1-schnell"
-revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-exclude = ["transformer"]
-
-[[tool.edge-maxxing.models]]
-repository = "RobertML/FLUX.1-schnell-int8wo"
-revision = "307e0777d92df966a3c0f99f31a6ee8957a9857a"
-
-[[tool.edge-maxxing.models]]
-repository = "city96/t5-v1_1-xxl-encoder-bf16"
-revision = "1b9c856aadb864af93c1dcdc226c2774fa67bc86"
-
-[[tool.edge-maxxing.models]]
-repository = "proact/PRO_FLUX_012"
-revision = "44b2be6e38d78508fc1e6bdf768fa8a4dbf61fcd"
-
+repository = "slobers/Flux.1.Schnella"
+revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
 
 [project.scripts]
-start_inference = "main:main"
-
+start_inference = "main:main"

src/main.py

@@ -1,24 +1,18 @@
-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")
 
+from pipeline import load_pipeline, infer
 
-def at_exit():
-    torch.cuda.empty_cache()
+SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")
 
 
 def main():
-    atexit.register(at_exit)
-
     print(f"Loading pipeline")
     pipeline = load_pipeline()
 
@@ -33,7 +27,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"))
@@ -41,13 +35,15 @@ def main():
                     print(f"Inference socket exiting")
 
                     return
-                image = infer(request, pipeline, generator.manual_seed(request.seed))
+
+                image = infer(request, pipeline)
+
                 data = BytesIO()
                 image.save(data, format=JpegImageFile.format)
 
                 packet = data.getvalue()
 
-                connection.send_bytes(packet )
+                connection.send_bytes(packet)
 
 
 if __name__ == '__main__':

src/pipeline.py

@@ -1,73 +1,46 @@
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
-
+#4
+#4
 from huggingface_hub.constants import HF_HUB_CACHE
-from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer
+from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
 import torch
 import torch._dynamo
 import gc
-from PIL import Image as img
+import os
+from diffusers import FluxPipeline, AutoencoderTiny
 from PIL.Image import Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
-import time
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
-import os
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
-
-
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.enabled = True
-torch.backends.cudnn.benchmark = True
-
-ckpt_id = "black-forest-labs/FLUX.1-schnell"
-ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
-
-
+os.environ["TOKENIZERS_PARALLELISM"] = "True"
+torch._dynamo.config.suppress_errors = True
 Pipeline = None
-def empty_cache():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_peak_memory_stats()
-
-def load_pipeline() -> Pipeline:    
-    empty_cache()
-
-    dtype, device = torch.bfloat16, "cuda"
-    
-    text_encoder_2 = T5EncoderModel.from_pretrained(
-        "city96/t5-v1_1-xxl-encoder-bf16", revision = "1b9c856aadb864af93c1dcdc226c2774fa67bc86", torch_dtype=torch.bfloat16
-    ).to(memory_format=torch.channels_last)
-    
-
-    path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
-    model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
-    pipeline = DiffusionPipeline.from_pretrained(
-        ckpt_id,
-        revision=ckpt_revision,
-        transformer=model,
-        text_encoder_2=text_encoder_2,
-        torch_dtype=dtype,
-        ).to(device)
+ids = "slobers/Flux.1.Schnella"
+Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
+gc.collect()
+torch.cuda.empty_cache()
+torch.cuda.reset_max_memory_allocated()
+torch.cuda.reset_peak_memory_stats()
+def load_pipeline() -> Pipeline:
+    path = os.path.join(HF_HUB_CACHE, "models--slobers--Flux.1.Schnella/snapshots/e34d670e44cecbbc90e4962e7aada2ac5ce8b55b/transformer")
+    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False).to(memory_format=torch.channels_last)
+    pipeline = FluxPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
+    pipeline.to("cuda")
     quantize_(pipeline.vae, int8_weight_only())
-
-    pipeline(prompt="imprisonable, forechamber, monotropic, blandiloquious, blechnoid", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
-    
-    empty_cache()
+    pipeline.transformer = torch.compile(pipeline.transformer, options={"shape_padding": True})
+    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
-
-
 @torch.no_grad()
-def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
-    
-    image=pipeline(request.prompt,
-                generator=generator, 
-                guidance_scale=0.0, 
-                num_inference_steps=4, 
-                max_sequence_length=256, 
-                height=request.height, 
-                width=request.width, 
-                output_type="pil").images[0]
-
-    return(image)
+def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
+    generator = Generator(pipeline.device).manual_seed(request.seed)
+    return pipeline(
+        request.prompt,
+        generator=generator,
+        guidance_scale=0.0,
+        num_inference_steps=4,
+        max_sequence_length=256,
+        height=request.height,
+        width=request.width,
+    ).images[0]

uv.lock

@@ -1,15 +1,8 @@
 version = 1
 requires-python = ">=3.10, <3.13"
 resolution-markers = [
-    "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')",
+    "python_full_version < '3.12'",
+    "python_full_version >= '3.12'",
 ]
 
 [[package]]
@@ -160,7 +153,6 @@ dependencies = [
     { name = "accelerate" },
     { name = "diffusers" },
     { name = "edge-maxxing-pipelines" },
-    { name = "gitpython" },
     { name = "hf-transfer" },
     { name = "omegaconf" },
     { name = "protobuf" },
@@ -175,7 +167,6 @@ 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" },
@@ -194,30 +185,6 @@ 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"
@@ -445,7 +412,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.1.0.70"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "nvidia-cublas-cu12" },
 ]
 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 },
@@ -456,7 +423,7 @@ name = "nvidia-cufft-cu12"
 version = "11.2.1.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 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 },
@@ -477,9 +444,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.6.1.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { 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')" },
+    { name = "nvidia-cublas-cu12" },
+    { name = "nvidia-cusparse-cu12" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 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 },
@@ -491,7 +458,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.3.1.170"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 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 },
@@ -888,15 +855,6 @@ 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"
@@ -1055,7 +1013,7 @@ name = "triton"
 version = "3.1.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "filelock", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux')" },
+    { name = "filelock" },
 ]
 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 },