Initial commit with folder contents

.gitattributes

@@ -32,6 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.xz 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
+*tfevents* filter=lfs diff=lfs merge=lfs -text

pyproject.toml

@@ -27,14 +27,13 @@ repository = "black-forest-labs/FLUX.1-schnell"
 revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
 exclude = ["transformer"]
 
-[[tool.edge-maxxing.models]]
-repository = "madebyollin/taef1"
-revision = "2d552378e58c9c94201075708d7de4e1163b2689"
-
 [[tool.edge-maxxing.models]]
 repository = "farapart/t5_encoder"
 revision = "c225a976e16b77764f653a801268de86e20adb84"
 
+[[tool.edge-maxxing.models]]
+repository = "madebyollin/taef1"
+revision = "2d552378e58c9c94201075708d7de4e1163b2689"
 
 [project.scripts]
 start_inference = "main:main"

src/pipeline.py

@@ -1,55 +1,26 @@
-from diffusers import (
-    DiffusionPipeline,
-    AutoencoderKL,
-    FluxPipeline,
-    FluxTransformer2DModel
-)
-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
+import os
+import torch.nn.functional as F
 from PIL import Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
-import time
-import math
-from typing import Type, Dict, Any, Tuple, Callable, Optional, Union
-import numpy as np
-import torch.nn as nn
-import torch.nn.functional as F
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+from typing import Type
+from diffusers import DiffusionPipeline, FluxTransformer2DModel
+from huggingface_hub.constants import HF_HUB_CACHE
+from transformers import T5EncoderModel
 
-# preconfigs
-import os
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
-# torch.backends.cudnn.benchmark = True
 
-# globals
 Pipeline = None
-ckpt_id = "black-forest-labs/FLUX.1-schnell"
-ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
-TinyVAE = "madebyollin/taef1"
-TinyVAE_REV = "2d552378e58c9c94201075708d7de4e1163b2689"
-
-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:
+    ckpt_id = "black-forest-labs/FLUX.1-schnell"
+    ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
     text_encoder_2 = T5EncoderModel.from_pretrained("farapart/t5_encoder", revision = "c225a976e16b77764f653a801268de86e20adb84", subfolder="text_encoder_2",torch_dtype=torch.bfloat16)
     path = os.path.join(HF_HUB_CACHE, "models--farapart--t5_encoder/snapshots/c225a976e16b77764f653a801268de86e20adb84/transformer")
     transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
@@ -60,11 +31,6 @@ def load_pipeline() -> Pipeline:
         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
 
-sample = 1
 @torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
-    global sample
-    if not sample:
-        sample=1
-        empty_cache()
     return 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]