Add files using upload-large-folder tool

src/pipeline.py

@@ -0,0 +1,76 @@
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+
+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
+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"
+
+Pipeline = None
+
+ckpt_id = "black-forest-labs/FLUX.1-schnell"
+def empty_cache():
+    start = time.time()
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    print(f"Flush took: {time.time() - start}")
+
+def load_pipeline() -> Pipeline:    
+    empty_cache()
+
+    dtype, device = torch.bfloat16, "cuda"
+
+    vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_e3m2", torch_dtype=dtype)
+    quantize_(vae, int8_weight_only())
+
+    ############ Text Encoder ############
+    text_encoder = CLIPTextModel.from_pretrained(
+        ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
+    )
+    quantize_(text_encoder, int8_weight_only())
+    ############ Text Encoder 2 ############
+    text_encoder_2 = T5EncoderModel.from_pretrained(
+        "city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16
+    )
+    quantize_(text_encoder_2, int8_weight_only())
+    model = FluxTransformer2DModel.from_pretrained(
+        "/home/sandbox/.cache/huggingface/hub/models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a", torch_dtype=dtype, use_safetensors=False
+    )
+    pipeline = DiffusionPipeline.from_pretrained(
+        ckpt_id, 
+        transformer=model,
+        text_encoder=text_encoder,
+        text_encoder_2=text_encoder_2,
+        torch_dtype=dtype,
+        vae=vae
+        ).to(device)
+
+    for _ in range(2):
+        pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
+    
+    empty_cache()
+    return pipeline
+
+
+from datetime import datetime
+@torch.inference_mode()
+def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
+    try:
+        generator = Generator(pipeline.device).manual_seed(request.seed)
+        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]
+    except:
+        image = img.open("./RobertML.png")
+        pass
+    return(image)