| import os |
| from diffusers import FluxPipeline, AutoencoderKL, FluxTransformer2DModel |
| from diffusers.image_processor import VaeImageProcessor |
| from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel, CLIPTextConfig, T5Config |
| import torch |
| import gc |
| from PIL import Image |
| from pipelines.models import TextToImageRequest |
| from torch import Generator |
| from time import perf_counter |
|
|
|
|
| HOME = os.environ["HOME"] |
| os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01" |
| os.environ['PYTHONMALLOC'] = 'malloc' |
|
|
| CHECKPOINT = "black-forest-labs/FLUX.1-schnell" |
|
|
| DTYPE = torch.bfloat16 |
| NUM_STEPS = 4 |
|
|
|
|
| def empty_cache(): |
| gc.collect() |
| torch.cuda.empty_cache() |
| torch.cuda.reset_max_memory_allocated() |
| torch.cuda.reset_peak_memory_stats() |
|
|
|
|
| def load_pipeline() -> FluxPipeline: |
| empty_cache() |
| |
| pipe = FluxPipeline.from_pretrained(CHECKPOINT, torch_dtype=DTYPE) |
| |
| pipe.text_encoder_2.to(memory_format=torch.channels_last) |
| pipe.transformer.to(memory_format=torch.channels_last) |
| |
| pipe.vae.to(memory_format=torch.channels_last) |
| pipe.vae = torch.compile(pipe.vae) |
| pipe._exclude_from_cpu_offload = ["vae"] |
| |
| pipe.enable_sequential_cpu_offload() |
| |
| empty_cache() |
| pipe("dog", guidance_scale=0.0, max_sequence_length=256, num_inference_steps=4) |
| return pipe |
|
|
| @torch.inference_mode() |
| def infer(request: TextToImageRequest, _pipeline: FluxPipeline) -> Image: |
| torch.cuda.reset_peak_memory_stats() |
| |
| if request.seed is None: |
| generator = None |
| else: |
| generator = Generator(device="cuda").manual_seed(request.seed) |
|
|
| empty_cache() |
| image = _pipeline(prompt=request.prompt, |
| width=request.width, |
| height=request.height, |
| guidance_scale=0.0, |
| generator=generator, |
| output_type="pil", |
| max_sequence_length=256, |
| num_inference_steps=NUM_STEPS).images[0] |
| return image |
|
|
