Initial commit

.gitattributes

@@ -0,0 +1,5 @@
+src/taef1_decoder_only.pth filter=lfs diff=lfs merge=lfs -text
+taef1_decoder_only.pth filter=lfs diff=lfs merge=lfs -text
+decoder.pth filter=lfs diff=lfs merge=lfs -text
+encoder.pth filter=lfs diff=lfs merge=lfs -text
+backup.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,19 @@
+# flux-schnell-edge-inference
+
+This holds the baseline for the FLUX Schnel NVIDIA GeForce RTX 4090 contest, which can be forked freely and optimized
+
+Some recommendations are as follows:
+- Installing dependencies should be done in `pyproject.toml`, including git dependencies
+- HuggingFace models should be specified in the `models` array in the `pyproject.toml` file, and will be downloaded before benchmarking
+    - The pipeline does **not** have internet access so all dependencies and models must be included in the `pyproject.toml`
+    - Compiled models should be hosted on HuggingFace and included in the `models` array in the `pyproject.toml` (rather than compiling during loading). Loading time matters far more than file sizes
+- Avoid changing `src/main.py`, as that includes mostly protocol logic. Most changes should be in `models` and `src/pipeline.py`
+- Ensure the entire repository (excluding dependencies and HuggingFace models) is under 16MB
+
+For testing, you need a docker container with pytorch and ubuntu 22.04.
+You can download your listed dependencies with `uv`, installed with:
+```bash
+pipx ensurepath
+pipx install uv
+```
+You can then relock with `uv lock`, and then run with `uv run start_inference`

decoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4142638173f27b872a916b2a6b599fe311f22fe3fd84c77079af5bc114928490
+size 1800500

encoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c10a3d71556ddc04705f276bcffeaed53072e1342502a4d121ed99f585640e60
+size 4940983

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools >= 75.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flux-schnell-edge-inference"
+description = "An edge-maxxing model submission for the 4090 Flux contest"
+requires-python = ">=3.10,<3.13"
+version = "7"
+dependencies = [
+  "diffusers==0.31.0",
+  "transformers==4.46.2",
+  "accelerate==1.1.0",
+  "omegaconf==2.3.0",
+  "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",
+  "torchao>=0.6.1", 
+  "torchvision"
+]
+
+[tool.edge-maxxing]
+models = ["black-forest-labs/FLUX.1-schnell", "madebyollin/taef1"]
+
+[project.scripts]
+start_inference = "main:main"

src/main.py

@@ -0,0 +1,59 @@
+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}'")
+
+    if exists(SOCKET):
+        remove(SOCKET)
+
+    with Listener(SOCKET) as listener:
+        chmod(SOCKET, 0o777)
+
+        print(f"Awaiting connections")
+        with listener.accept() as connection:
+            print(f"Connected")
+
+            while True:
+                try:
+                    request = TextToImageRequest.model_validate_json(connection.recv_bytes().decode("utf-8"))
+                except EOFError:
+                    print(f"Inference socket exiting")
+
+                    return
+
+                image = infer(request, pipeline)
+
+                data = BytesIO()
+                image.save(data, format=JpegImageFile.format)
+
+                packet = data.getvalue()
+
+                connection.send_bytes(packet)
+
+
+if __name__ == '__main__':
+    main()

src/model.py

@@ -0,0 +1,137 @@
+
+import torch
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+class Block(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.fuse = nn.ReLU()
+    def forward(self, x):
+        return self.fuse(self.conv(x) + self.skip(x))
+
+def Encoder(latent_channels=4):
+    return nn.Sequential(
+        conv(3, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, latent_channels),
+    )
+
+class DCAH(nn.Module):
+    def __init__(self, in_channels, embed_dim=64, dilation_rates=(1, 2, 4)):
+        super(DCAH, self).__init__()
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+        self.dilated_convs = nn.ModuleList([
+            nn.Conv2d(in_channels, embed_dim, kernel_size=3, padding=rate, dilation=rate)
+            for rate in dilation_rates
+        ])
+        self.dilated_conv_merge = nn.Conv2d(embed_dim * len(dilation_rates), embed_dim, kernel_size=1)
+        self.query = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
+        self.key = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
+        self.value = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
+        self.refine = nn.Sequential(
+            nn.Conv2d(embed_dim, embed_dim, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(embed_dim, in_channels, kernel_size=1)
+        )
+
+    def forward(self, x):
+        dilated_features = [conv(x) for conv in self.dilated_convs]
+        concat_features = torch.cat(dilated_features, dim=1)
+        global_context = self.dilated_conv_merge(concat_features)
+        q = self.query(global_context)
+        k = self.key(global_context)
+        v = self.value(global_context)
+        attention = F.softmax(torch.matmul(q.flatten(2), k.flatten(2).transpose(-2, -1)), dim=-1)
+        attention_out = torch.matmul(attention, v.flatten(2)).view_as(global_context)
+        refined = self.refine(global_context + attention_out)
+        return refined
+
+def DecoderSeq(latent_channels=16): 
+    return nn.Sequential(
+        Clamp(), 
+        conv(latent_channels, 48), 
+        nn.ReLU(),
+        Block(48, 48), Block(48, 48),  
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48), Block(48, 48), 
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48), 
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48), 
+        conv(48, 3), 
+    )
+
+
+class Decoder(nn.Module):
+    def __init__(self, latent_channels=16):
+        decoder = DecoderSeq(latent_channels=latent_channels)
+        refinement_head = DCAH(in_channels=3, embed_dim=64)
+        super(Decoder, self).__init__()
+        self.decoder = decoder
+        self.refinement_head = refinement_head
+
+    def forward(self, x):
+        decoded = self.decoder(x)
+        refined = self.refinement_head(decoded)
+        return refined
+        
+
+class Model(nn.Module):
+    latent_magnitude = 3
+    latent_shift = 0.5
+
+    def __init__(self, encoder_path="encoder.pth", decoder_path="decoder.pth", latent_channels=None):
+        super().__init__()
+        if latent_channels is None:
+            latent_channels = self.guess_latent_channels(str(encoder_path))
+        self.encoder = Encoder(latent_channels)
+        self.decoder = Decoder(latent_channels)
+        if encoder_path is not None:
+            encoder_state_dict = torch.load(encoder_path, map_location="cpu", weights_only=True)
+            filtered_state_dict = {k.strip('encoder.'): v for k, v in encoder_state_dict.items() if k.strip('encoder.') in self.encoder.state_dict() and v.size() == self.encoder.state_dict()[k.strip('encoder.')].size()}
+            self.encoder.load_state_dict(filtered_state_dict, strict=False)
+            
+        if decoder_path is not None:
+            decoder_state_dict = torch.load(decoder_path, map_location="cpu", weights_only=True)
+            filtered_state_dict = {k: v for k, v in decoder_state_dict.items() if k in self.decoder.state_dict() and v.size() == self.decoder.state_dict()[k].size()}
+            self.decoder.load_state_dict(filtered_state_dict, strict=False)
+            
+        self.encoder.requires_grad_(False)
+        self.decoder.decoder.requires_grad_(False)
+        self.decoder.refinement_head.requires_grad_(False)
+
+    def guess_latent_channels(self, encoder_path):
+        if "taef1" in encoder_path:return 16
+        if "taesd3" in encoder_path:return 16
+        return 4
+        
+    @staticmethod
+    def scale_latents(x):
+        return x.div(2 * Model.latent_magnitude).add(Model.latent_shift).clamp(0, 1)
+
+    @staticmethod
+    def unscale_latents(x):
+        return x.sub(Model.latent_shift).mul(2 * Model.latent_magnitude)
+
+    def forward(self, x, return_latent=False):
+        latent = self.encoder(x)
+        out = self.decoder(latent)
+        if return_latent:
+            return out.clamp(0, 1), latent
+        return out.clamp(0, 1)
+
+        

src/pipeline.py

@@ -0,0 +1,79 @@
+from diffusers import AutoencoderKL, AutoencoderTiny
+from diffusers.image_processor import VaeImageProcessor
+import torch
+import torch._dynamo
+import gc
+from PIL.Image import Image
+from pipelines.models import TextToImageRequest
+from torch import Generator
+from diffusers import FluxPipeline
+from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+import torch.nn as nn 
+from model import Model, Decoder
+import torchvision
+
+Pipeline = None
+MODEL_ID = "black-forest-labs/FLUX.1-schnell"
+DTYPE = torch.bfloat16
+def clear():
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+def load_pipeline() -> Pipeline:    
+    clear()
+    
+    # vae = Model("encoder.pth", "decoder.pth")
+    # vae.to(dtype=DTYPE)
+
+    vae = AutoencoderTiny.from_pretrained("madebyollin/taef1")
+    vae.decoder = Decoder(16)
+    decoder_path = "decoder.pth"
+    decoder_state_dict = torch.load(decoder_path,  weights_only=True) #map_location="cpu",
+    filtered_state_dict = {k: v for k, v in decoder_state_dict.items() if k in vae.decoder.state_dict() and v.size() == vae.decoder.state_dict()[k].size()}
+    print(decoder_state_dict.keys())
+    print(filtered_state_dict.keys())
+    vae.decoder.load_state_dict(filtered_state_dict, strict=False)
+    vae.decoder.requires_grad_(False)
+    vae.to(dtype=DTYPE)
+    
+    # quantize_(vae, fpx_weight_only(3, 2))
+    # quantize_(vae, int8_weight_only())
+    pipeline = FluxPipeline.from_pretrained(MODEL_ID,vae=vae,
+                                        torch_dtype=DTYPE)
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.cuda.set_per_process_memory_fraction(0.99)
+    pipeline.text_encoder.to(memory_format=torch.channels_last)
+    pipeline.text_encoder_2.to(memory_format=torch.channels_last)
+    pipeline.transformer.to(memory_format=torch.channels_last)
+    pipeline.vae.to(memory_format=torch.channels_last)
+    pipeline.vae = torch.compile(pipeline.vae)
+    pipeline._exclude_from_cpu_offload = ["vae"]
+    pipeline.enable_sequential_cpu_offload()
+    torch.jit.enable_onednn_fusion(True)
+    clear()
+    for _ in range(1):
+        pipeline(prompt="unpervaded, unencumber, froggish, groundneedle, transnatural, fatherhood, outjump, cinerator", width=1024, height=1024, guidance_scale=0.1, num_inference_steps=4, max_sequence_length=256)
+    return pipeline
+
+sample = True
+@torch.inference_mode()
+def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
+    global sample
+    if sample:
+        clear()
+        sample = None
+    torch.cuda.reset_peak_memory_stats()
+    generator = Generator("cuda").manual_seed(request.seed)
+    image = None
+    # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+    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="pt").images[0]
+    # image = image / 255.
+    # image = image.mul_(2).sub_(1)
+    # image = ((image + 1) / 2) * 255
+    # image = image.clamp(0, 255)
+    # image = image.to(torch.float32)
+    # return torchvision.transforms.functional.to_pil_image(image)
+    return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))