Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ 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
+backup.png filter=lfs diff=lfs merge=lfs -text

pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools >= 75.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flux-schnell-edge-inference"
+description = "An edge-maxxing model submission by RobertML for the 4090 Flux contest"
+requires-python = ">=3.10,<3.13"
+version = "8"
+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",
+  "hf_transfer==0.1.8",
+  "torchao==0.6.1",
+]
+
+[[tool.edge-maxxing.models]]
+repository = "madebyollin/taef1"
+revision = "2d552378e58c9c94201075708d7de4e1163b2689"
+
+[[tool.edge-maxxing.models]]
+repository = "director432/Flux1-Schnell"
+revision = "f8da0eb2b421c7677c70312bd9dec91a71f411d3"
+
+[[tool.edge-maxxing.models]]
+repository = "director432/Flux1-Transformer2D"
+revision = "803893c49df2bb29d3f0f89ef5467781bef64b25"
+
+[[tool.edge-maxxing.models]]
+repository = "director432/Flux1-T5Encoder"
+revision = "93fa999c3acb891488a05eebfb6a98e31d574d05"
+
+
+[project.scripts]
+start_inference = "main:main"
+

src/flux_schnell_edge_inference.egg-info/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.1
+Name: flux-schnell-edge-inference
+Version: 7
+Summary: An edge-maxxing model submission for the 4090 Flux contest
+Requires-Python: <3.13,>=3.10
+Requires-Dist: diffusers==0.31.0
+Requires-Dist: transformers==4.46.2
+Requires-Dist: accelerate==1.1.0
+Requires-Dist: omegaconf==2.3.0
+Requires-Dist: torch==2.5.1
+Requires-Dist: protobuf==5.28.3
+Requires-Dist: sentencepiece==0.2.0
+Requires-Dist: edge-maxxing-pipelines@ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines
+Requires-Dist: gitpython>=3.1.43
+Requires-Dist: torchao>=0.6.1

src/flux_schnell_edge_inference.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+README.md
+pyproject.toml
+src/main.py
+src/pipeline.py
+src/flux_schnell_edge_inference.egg-info/PKG-INFO
+src/flux_schnell_edge_inference.egg-info/SOURCES.txt
+src/flux_schnell_edge_inference.egg-info/dependency_links.txt
+src/flux_schnell_edge_inference.egg-info/entry_points.txt
+src/flux_schnell_edge_inference.egg-info/requires.txt
+src/flux_schnell_edge_inference.egg-info/top_level.txt

src/flux_schnell_edge_inference.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/flux_schnell_edge_inference.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+start_inference = main:main

src/flux_schnell_edge_inference.egg-info/requires.txt

@@ -0,0 +1,10 @@
+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

src/flux_schnell_edge_inference.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+main
+pipeline

src/ghanta.py

@@ -0,0 +1,74 @@
+import torch
+from typing import Tuple, Callable
+def hacer_nada(x: torch.Tensor, modo: str = None):
+    return x
+def brujeria_mps(entrada, dim, indice):
+    if entrada.shape[-1] == 1:
+        return torch.gather(entrada.unsqueeze(-1), dim - 1 if dim < 0 else dim, indice.unsqueeze(-1)).squeeze(-1)
+    else:
+        return torch.gather(entrada, dim, indice)
+def emparejamiento_suave_aleatorio_2d(
+    metrica: torch.Tensor,
+    ancho: int,
+    alto: int,
+    paso_x: int,
+    paso_y: int,
+    radio: int,
+    sin_aleatoriedad: bool = False,
+    generador: torch.Generator = None
+) -> Tuple[Callable, Callable]:
+    lote, num_nodos, _ = metrica.shape
+    if radio <= 0:
+        return hacer_nada, hacer_nada
+    recopilar = brujeria_mps if metrica.device.type == "mps" else torch.gather
+    with torch.no_grad():
+        alto_paso_y, ancho_paso_x = alto // paso_y, ancho // paso_x
+        if sin_aleatoriedad:
+            indice_aleatorio = torch.zeros(alto_paso_y, ancho_paso_x, 1, device=metrica.device, dtype=torch.int64)
+        else:
+            indice_aleatorio = torch.randint(paso_y * paso_x, size=(alto_paso_y, ancho_paso_x, 1), device=generador.device, generator=generador).to(metrica.device)
+        vista_buffer_indice = torch.zeros(alto_paso_y, ancho_paso_x, paso_y * paso_x, device=metrica.device, dtype=torch.int64)
+        vista_buffer_indice.scatter_(dim=2, index=indice_aleatorio, src=-torch.ones_like(indice_aleatorio, dtype=indice_aleatorio.dtype))
+        vista_buffer_indice = vista_buffer_indice.view(alto_paso_y, ancho_paso_x, paso_y, paso_x).transpose(1, 2).reshape(alto_paso_y * paso_y, ancho_paso_x * paso_x)
+        if (alto_paso_y * paso_y) < alto or (ancho_paso_x * paso_x) < ancho:
+            buffer_indice = torch.zeros(alto, ancho, device=metrica.device, dtype=torch.int64)
+            buffer_indice[:(alto_paso_y * paso_y), :(ancho_paso_x * paso_x)] = vista_buffer_indice
+        else:
+            buffer_indice = vista_buffer_indice
+        indice_aleatorio = buffer_indice.reshape(1, -1, 1).argsort(dim=1)
+        del buffer_indice, vista_buffer_indice
+        num_destino = alto_paso_y * ancho_paso_x
+        indices_a = indice_aleatorio[:, num_destino:, :]
+        indices_b = indice_aleatorio[:, :num_destino, :]
+        def dividir(x):
+            canales = x.shape[-1]
+            origen = recopilar(x, dim=1, index=indices_a.expand(lote, num_nodos - num_destino, canales))
+            destino = recopilar(x, dim=1, index=indices_b.expand(lote, num_destino, canales))
+            return origen, destino
+        metrica = metrica / metrica.norm(dim=-1, keepdim=True)
+        a, b = dividir(metrica)
+        puntuaciones = a @ b.transpose(-1, -2)
+        radio = min(a.shape[1], radio)
+        nodo_max, nodo_indice = puntuaciones.max(dim=-1)
+        indice_borde = nodo_max.argsort(dim=-1, descending=True)[..., None]
+        indice_no_emparejado = indice_borde[..., radio:, :]
+        indice_origen = indice_borde[..., :radio, :]
+        indice_destino = recopilar(nodo_indice[..., None], dim=-2, index=indice_origen)
+    def fusionar(x: torch.Tensor, modo="mean") -> torch.Tensor:
+        origen, destino = dividir(x)
+        n, t1, c = origen.shape
+        no_emparejado = recopilar(origen, dim=-2, index=indice_no_emparejado.expand(n, t1 - radio, c))
+        origen = recopilar(origen, dim=-2, index=indice_origen.expand(n, radio, c))
+        destino = destino.scatter_reduce(-2, indice_destino.expand(n, radio, c), origen, reduce=modo)
+        return torch.cat([no_emparejado, destino], dim=1)
+    def desfusionar(x: torch.Tensor) -> torch.Tensor:
+        longitud_no_emparejado = indice_no_emparejado.shape[1]
+        no_emparejado, destino = x[..., :longitud_no_emparejado, :], x[..., longitud_no_emparejado:, :]
+        _, _, c = no_emparejado.shape
+        origen = recopilar(destino, dim=-2, index=indice_destino.expand(lote, radio, c))
+        salida = torch.zeros(lote, num_nodos, c, device=x.device, dtype=x.dtype)
+        salida.scatter_(dim=-2, index=indices_b.expand(lote, num_destino, c), src=destino)
+        salida.scatter_(dim=-2, index=recopilar(indices_a.expand(lote, indices_a.shape[1], 1), dim=1, index=indice_no_emparejado).expand(lote, longitud_no_emparejado, c), src=no_emparejado)
+        salida.scatter_(dim=-2, index=recopilar(indices_a.expand(lote, indices_a.shape[1], 1), dim=1, index=indice_origen).expand(lote, radio, c), src=origen)
+        return salida
+    return fusionar, desfusionar

src/main.py

@@ -0,0 +1,55 @@
+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
+from git import Repo
+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")
+            generator = torch.Generator("cuda")
+            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, generator.manual_seed(request.seed))
+                data = BytesIO()
+                image.save(data, format=JpegImageFile.format)
+
+                packet = data.getvalue()
+
+                connection.send_bytes(packet )
+
+
+if __name__ == '__main__':
+    main()

src/pipeline.py

@@ -0,0 +1,511 @@
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from huggingface_hub.constants import HF_HUB_CACHE
+from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
+import torch
+import torch._dynamo
+import gc
+from PIL import Image as img
+from PIL.Image import Image
+from pipelines.models import TextToImageRequest
+from torch import Generator
+import time
+import os
+
+import torch
+import math
+from typing import Type, Dict, Any, Tuple, Callable, Optional, Union, List
+import ghanta 
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    FluxAttnProcessor2_0,
+    FusedFluxAttnProcessor2_0,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+from diffusers.loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+# from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+
+torch._dynamo.config.suppress_errors = True
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+dtype = torch.bfloat16
+device = "cuda"
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+def calc_shift(img_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.16):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    return img_seq_len * m + b
+
+def get_timesteps(scheduler, num_steps: Optional[int] = None, dev: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, **kwargs):
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=dev, **kwargs)
+        num_steps = len(scheduler.timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=dev, **kwargs)
+        num_steps = len(scheduler.timesteps)
+    else:
+        scheduler.set_timesteps(num_steps, device=dev, **kwargs)
+    return scheduler.timesteps, num_steps
+
+def init_generator(dev: torch.device, backup: Optional[torch.Generator] = None):
+    if dev.type == "cpu":
+        return torch.Generator(device="cpu").set_state(torch.get_rng_state())
+    elif dev.type == "cuda":
+        return torch.Generator(device=dev).set_state(torch.cuda.get_rng_state())
+    else:
+        return init_generator(torch.device("cpu")) if backup is None else backup
+
+def compute_fusion(x: torch.Tensor, info: Dict[str, Any]) -> Tuple[Callable, ...]:
+    h_orig, w_orig = info["size"]
+    tokens = h_orig * w_orig
+    sub_sample = int(math.ceil(math.sqrt(tokens // x.shape[1])))
+    args = info["args"]
+    if sub_sample <= args["down"]:
+        new_w = int(math.ceil(w_orig / sub_sample))
+        new_h = int(math.ceil(h_orig / sub_sample))
+        ratio = int(x.shape[1] * args["ratio"])
+        if args["generator"] is None:
+            args["generator"] = init_generator(x.device)
+        elif args["generator"].device != x.device:
+            args["generator"] = init_generator(x.device, backup=args["generator"])
+        use_rand = args["rando"]
+        fusion, defusion = ghanta.emparejamiento_suave_aleatorio_2d(x, new_w, new_h, args["sx"], args["sy"], ratio, sin_aleatoriedad=not use_rand, generador=args["generator"])
+    else:
+        fusion, defusion = (ghanta.hacer_nada, ghanta.hacer_nada)
+    fusion_a, defusion_a = (fusion, defusion) if args["m1"] else (ghanta.hacer_nada, ghanta.hacer_nada)
+    fusion_c, defusion_c = (fusion, defusion) if args["m2"] else (ghanta.hacer_nada, ghanta.hacer_nada)
+    fusion_m, defusion_m = (fusion, defusion) if args["m3"] else (ghanta.hacer_nada, ghanta.hacer_nada)
+    return fusion_a, fusion_c, fusion_m, defusion_a, defusion_c, defusion_m
+
+@maybe_allow_in_graph
+class SingleTransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden = int(dim * mlp_ratio)
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.mlp_proj = nn.Linear(dim, self.mlp_hidden)
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.out_proj = nn.Linear(dim + self.mlp_hidden, dim)
+        proc = FluxAttnProcessor2_0()
+        self.attn = Attention(query_dim=dim, cross_attention_dim=None, dim_head=head_dim, heads=num_heads, out_dim=dim, bias=True, processor=proc, qk_norm="rms_norm", eps=1e-6, pre_only=True)
+    def forward(self, hidden: torch.FloatTensor, temb: torch.FloatTensor, image_rotary_emb=None, joint_attention_kwargs=None, tinfo: Dict[str, Any] = None):
+        if tinfo is not None:
+            m_a, m_c, mom, u_a, u_c, u_m = compute_fusion(hidden, tinfo)
+        else:
+            m_a, m_c, mom, u_a, u_c, u_m = (ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada)
+        residual = hidden
+        norm_hidden, gate = self.norm(hidden, emb=temb)
+        norm_hidden = m_a(norm_hidden)
+        mlp_hidden = self.mlp_act(self.mlp_proj(norm_hidden))
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_out = self.attn(hidden_states=norm_hidden, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs)
+        hidden = torch.cat([attn_out, mlp_hidden], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden = gate * self.out_proj(hidden)
+        hidden = u_a(residual + hidden)
+        return hidden
+
+@maybe_allow_in_graph
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, head_dim, qk_norm="rms_norm", eps=1e-6):
+        super().__init__()
+        self.norm1 = AdaLayerNormZero(dim)
+        self.norm1_context = AdaLayerNormZero(dim)
+        if hasattr(F, "scaled_dot_product_attention"):
+            proc = FluxAttnProcessor2_0()
+        self.attn = Attention(query_dim=dim, cross_attention_dim=None, added_kv_proj_dim=dim, dim_head=head_dim, heads=num_heads, out_dim=dim, context_pre_only=False, bias=True, processor=proc, qk_norm=qk_norm, eps=eps)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        self._chunk_size = None
+        self._chunk_dim = 0
+    def forward(self, hidden: torch.FloatTensor, encoder_hidden: torch.FloatTensor, temb: torch.FloatTensor, image_rotary_emb=None, joint_attention_kwargs=None, tinfo: Dict[str, Any] = None):
+        if tinfo is not None:
+            m_a, m_c, mom, u_a, u_c, u_m = compute_fusion(hidden, tinfo)
+        else:
+            m_a, m_c, mom, u_a, u_c, u_m = (ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada, ghanta.hacer_nada)
+        norm_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden, emb=temb)
+        norm_encoder, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden, emb=temb)
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        norm_hidden = m_a(norm_hidden)
+        norm_encoder = m_c(norm_encoder)
+        attn_out, ctx_attn_out = self.attn(hidden_states=norm_hidden, encoder_hidden_states=norm_encoder, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs)
+        attn_out = gate_msa.unsqueeze(1) * attn_out
+        hidden = u_a(attn_out) + hidden
+        norm_hidden = self.norm2(hidden)
+        norm_hidden = norm_hidden * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_hidden = mom(norm_hidden)
+        ff_out = self.ff(norm_hidden)
+        ff_out = gate_mlp.unsqueeze(1) * ff_out
+        hidden = u_m(ff_out) + hidden
+        ctx_attn_out = c_gate_msa.unsqueeze(1) * ctx_attn_out
+        encoder_hidden = u_c(ctx_attn_out) + encoder_hidden
+        norm_encoder = self.norm2_context(encoder_hidden)
+        norm_encoder = norm_encoder * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        ctx_ff_out = self.ff_context(norm_encoder)
+        encoder_hidden = encoder_hidden + c_gate_mlp.unsqueeze(1) * ctx_ff_out
+        return encoder_hidden, hidden
+
+class Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
+    @register_to_config
+    def __init__(self, patch_size: int = 1, in_channels: int = 64, out_channels: Optional[int] = None, guidance_embeds: bool = False, axes_dims_rope: Tuple[int] = (16, 56, 56), generator: Optional[torch.Generator] = None):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+        text_time_cls = CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        self.time_text_embed = text_time_cls(embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim)
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(self.config.in_channels, self.inner_dim)
+        self.transformer_blocks = nn.ModuleList([TransformerBlock(dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim) for _ in range(self.config.num_layers)])
+        self.single_transformer_blocks = nn.ModuleList([SingleTransformerBlock(dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim) for _ in range(self.config.num_single_layers)])
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+        ratio: float = 0.5
+        down: int = 1
+        sx: int = 2
+        sy: int = 2
+        rando: bool = False
+        m1: bool = True
+        m2: bool = True
+        m3: bool = False
+        self.tinfo = {"size": None, "args": {"ratio": ratio, "down": down, "sx": sx, "sy": sy, "rando": rando, "m1": m1, "m2": m2, "m3": m3, "generator": generator}}
+        self.gradient_checkpointing = False
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        procs = {}
+        def add_processors(name: str, module: torch.nn.Module, procs: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                procs[f"{name}.processor"] = module.get_processor()
+            for sub_name, child in module.named_children():
+                add_processors(f"{name}.{sub_name}", child, procs)
+            return procs
+        for name, module in self.named_children():
+            add_processors(name, module, procs)
+        return procs
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        cnt = len(self.attn_processors.keys())
+        def set_proc(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+            for sub_name, child in module.named_children():
+                set_proc(f"{name}.{sub_name}", child, processor)
+        for name, module in self.named_children():
+            set_proc(name, module, processor)
+    def fuse_qkv_projections(self):
+        self.original_attn_processors = None
+        self.original_attn_processors = self.attn_processors
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+        self.set_attn_processor(FusedFluxAttnProcessor2_0())
+    def unfuse_qkv_projections(self):
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+    def forward(self, hidden: torch.Tensor, encoder_hidden: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_block_samples=None, controlnet_single_block_samples=None, return_dict: bool = True, controlnet_blocks_repeat: bool = False) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        if len(hidden.shape) == 4:
+            self.tinfo["size"] = (hidden.shape[2], hidden.shape[3])
+        if len(hidden.shape) == 3:
+            self.tinfo["size"] = (hidden.shape[1], hidden.shape[2])
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+        if USE_PEFT_BACKEND:
+            scale_lora_layers(self, lora_scale)
+        hidden = self.x_embedder(hidden)
+        timestep = timestep.to(hidden.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden.dtype) * 1000
+        else:
+            guidance = None
+        temb = self.time_text_embed(timestep, pooled_projections) if guidance is None else self.time_text_embed(timestep, guidance, pooled_projections)
+        encoder_hidden = self.context_embedder(encoder_hidden)
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+        for index, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                def custom_forward(module, return_dict=None):
+                    def cf(*inputs):
+                        return module(*inputs, return_dict=return_dict) if return_dict is not None else module(*inputs)
+                    return cf
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden, hidden = torch.utils.checkpoint.checkpoint(custom_forward(block), hidden, encoder_hidden, temb, image_rotary_emb, **ckpt_kwargs)
+            else:
+                encoder_hidden, hidden = block(hidden_states=hidden, encoder_hidden_states=encoder_hidden, temb=temb, image_rotary_emb=image_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, tinfo=self.tinfo)
+            if controlnet_block_samples is not None:
+                interval = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval = int(np.ceil(interval))
+                if controlnet_blocks_repeat:
+                    hidden = hidden + controlnet_block_samples[index % len(controlnet_block_samples)]
+                else:
+                    hidden = hidden + controlnet_block_samples[index // interval]
+        hidden = torch.cat([encoder_hidden, hidden], dim=1)
+        for index, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                def custom_forward(module, return_dict=None):
+                    def cf(*inputs):
+                        return module(*inputs, return_dict=return_dict) if return_dict is not None else module(*inputs)
+                    return cf
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden = torch.utils.checkpoint.checkpoint(custom_forward(block), hidden, temb, image_rotary_emb, **ckpt_kwargs)
+            else:
+                hidden = block(hidden_states=hidden, temb=temb, image_rotary_emb=image_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, tinfo=self.tinfo)
+            if controlnet_single_block_samples is not None:
+                interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval = int(np.ceil(interval))
+                hidden[:, encoder_hidden.shape[1]:, ...] = hidden[:, encoder_hidden.shape[1]:, ...] + controlnet_single_block_samples[index // interval]
+        hidden = hidden[:, encoder_hidden.shape[1]:, ...]
+        hidden = self.norm_out(hidden, temb)
+        output = self.proj_out(hidden)
+        if USE_PEFT_BACKEND:
+            unscale_lora_layers(self, lora_scale)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+class FluxPipeLine(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin):
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+    def __init__(self, scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, text_encoder_2: T5EncoderModel, tokenizer_2: T5TokenizerFast, transformer: Transformer2DModel):
+        super().__init__()
+        self.register_modules(vae=vae, text_encoder=text_encoder, text_encoder_2=text_encoder_2, tokenizer=tokenizer, tokenizer_2=tokenizer_2, transformer=transformer, scheduler=scheduler)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        self.default_sample_size = 128
+    def _get_t5_prompt_embeds(self, prompt: Union[str, List[str]] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, dev: Optional[torch.device] = None, dt: Optional[torch.dtype] = None):
+        dev = dev or self._execution_device
+        dt = dt or self.text_encoder.dtype
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2)
+        text_inputs = self.tokenizer_2(prompt, padding="max_length", max_length=max_sequence_length, truncation=True, return_length=False, return_overflowing_tokens=False, return_tensors="pt")
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(dev), output_hidden_states=False)[0]
+        dt = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dt, device=dev)
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        return prompt_embeds
+    def _get_clip_prompt_embeds(self, prompt: Union[str, List[str]], num_images_per_prompt: int = 1, dev: Optional[torch.device] = None):
+        dev = dev or self._execution_device
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+        text_inputs = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer_max_length, truncation=True, return_overflowing_tokens=False, return_length=False, return_tensors="pt")
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+        prompt_embeds = self.text_encoder(text_input_ids.to(dev), output_hidden_states=False)
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=dev)
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+        return prompt_embeds
+    def encode_prompt(self, prompt: Union[str, List[str]], prompt_2: Union[str, List[str]], dev: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, max_sequence_length: int = 512, lora_scale: Optional[float] = None):
+        dev = dev or self._execution_device
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(prompt=prompt, dev=dev, num_images_per_prompt=num_images_per_prompt)
+            prompt_embeds = self._get_t5_prompt_embeds(prompt=prompt_2, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, dev=dev)
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                unscale_lora_layers(self.text_encoder, lora_scale)
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+        dt = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=dev, dtype=dt)
+        return prompt_embeds, pooled_prompt_embeds, text_ids   
+    @staticmethod
+    def _prepare_latent_ids(batch_size, height, width, dev, dt):
+        latent_ids = torch.zeros(height, width, 3)
+        latent_ids[..., 1] = latent_ids[..., 1] + torch.arange(height)[:, None]
+        latent_ids[..., 2] = latent_ids[..., 2] + torch.arange(width)[None, :]
+        h, w, c = latent_ids.shape
+        latent_ids = latent_ids.reshape(h * w, c)
+        return latent_ids.to(device=dev, dtype=dt)
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels, height, width):
+        latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4)
+        return latents
+    @staticmethod
+    def _unpack_latents(latents, height, width, scale_factor):
+        batch_size, num_patches, channels = latents.shape
+        height = 2 * (int(height) // (scale_factor * 2))
+        width = 2 * (int(width) // (scale_factor * 2))
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        latents = latents.reshape(batch_size, channels // 4, height, width)
+        return latents
+    def prepare_latents(self, batch_size, num_channels, height, width, dt, dev, generator, latents=None):
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+        shape = (batch_size, num_channels, height, width)
+        if latents is not None:
+            latent_ids = self._prepare_latent_ids(batch_size, height // 2, width // 2, dev, dt)
+            return latents.to(device=dev, dtype=dt), latent_ids
+        latents = randn_tensor(shape, generator=generator, device=dev, dtype=dt)
+        latents = self._pack_latents(latents, batch_size, num_channels, height, width)
+        latent_ids = self._prepare_latent_ids(batch_size, height // 2, width // 2, dev, dt)
+        return latents, latent_ids
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+    @property
+    def interrupt(self):
+        return self._interrupt
+    @torch.no_grad()
+    def __call__(self, prompt: Union[str, List[str]] = None, prompt_2: Optional[Union[str, List[str]]] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 28, sigmas: Optional[List[float]] = None, guidance_scale: float = 3.5, num_images_per_prompt: Optional[int] = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, output_type: Optional[str] = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], max_sequence_length: int = 512):
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        dev = self._execution_device
+        lora_scale = self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(prompt=prompt, prompt_2=prompt_2, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, dev=dev, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, lora_scale=lora_scale)
+        num_channels = self.transformer.config.in_channels // 4
+        latents, latent_ids = self.prepare_latents(batch_size * num_images_per_prompt, num_channels, height, width, prompt_embeds.dtype, dev, generator, latents)
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calc_shift(image_seq_len, self.scheduler.config.base_image_seq_len, self.scheduler.config.max_image_seq_len, self.scheduler.config.base_shift, self.scheduler.config.max_shift)
+        timesteps, num_inference_steps = get_timesteps(self.scheduler, num_inference_steps, dev, sigmas=sigmas, mu=mu)
+        self._num_timesteps = len(timesteps)
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=dev, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            noise_pred = self.transformer(hidden_states=latents, timestep=timestep, guidance=guidance, pooled_projections=pooled_prompt_embeds, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False)[0]
+            latents_dtype = latents.dtype
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+            if latents.dtype != latents_dtype:
+                if torch.backends.mps.is_available():
+                    latents = latents.to(latents_dtype)
+            if callback_on_step_end is not None:
+                callback_kwargs = {}
+                for k in callback_on_step_end_tensor_inputs:
+                    callback_kwargs[k] = locals()[k]
+                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                latents = callback_outputs.pop("latents", latents)
+                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+            if XLA_AVAILABLE:
+                xm.mark_step()
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents / self.vae.config.scaling_factor + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (image,)
+        return FluxPipelineOutput(images=image)
+
+Pipeline = None
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+ckpt_id = "director432/Flux1-Schnell"
+ckpt_rev = "f8da0eb2b421c7677c70312bd9dec91a71f411d3"
+
+def clear_cache():
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+def load_pipeline() -> FluxPipeLine:
+    clear_cache()
+    text_enc2 = T5EncoderModel.from_pretrained("director432/Flux1-T5Encoder", revision="93fa999c3acb891488a05eebfb6a98e31d574d05", torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
+    path = os.path.join(HF_HUB_CACHE, "models--director432--Flux1-Transformer2D/snapshots/803893c49df2bb29d3f0f89ef5467781bef64b25")
+    gen = torch.Generator(device=device)
+    model = Transformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False, generator=gen).to(memory_format=torch.channels_last)
+    pipe = FluxPipeLine.from_pretrained(ckpt_id, revision=ckpt_rev, transformer=model, text_encoder_2=text_enc2, torch_dtype=dtype).to(device)
+    for _ in range(3):
+        pipe(prompt="director cooper follow interests activity recent organize", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
+    clear_cache()
+    return pipe
+
+@torch.no_grad()
+def infer(request: TextToImageRequest, pipe: FluxPipeLine, gen: Generator) -> Image:
+    return pipe(request.prompt, generator=gen, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]