src/pipeline.py

from diffusers import AutoencoderKL, FluxPipeline, FluxTransformer2DModel
from huggingface_hub.constants import HF_HUB_CACHE
from transformers import T5EncoderModel
from PIL import Image
from pipelines.models import TextToImageRequest
from torch import Generator
from typing import Type
from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
import torch
import torch._dynamo
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

Pipeline = None


def load_pipeline() -> Pipeline:
    ckpt_id = "intensity809/FLUX-encoder-VAE"
    ckpt_revision = "4dc7cf59f64cedde94422acff2ea8ef2b91a37d4"
    vae = AutoencoderKL.from_pretrained(
        ckpt_id,
        revision=ckpt_revision,
        subfolder="vae",
        local_files_only=True,
        torch_dtype=torch.bfloat16,
    )
    quantize_(vae, int8_weight_only())
    text_encoder_2 = T5EncoderModel.from_pretrained(
        "intensity809/FLUX1-transformer",
        revision="390b27551542e3b39c89ea96267c3a11f2e00980",
        subfolder="text_encoder_2",
        torch_dtype=torch.bfloat16,
    )
    path = os.path.join(
        HF_HUB_CACHE,
        "models--intensity809--FLUX1-transformer/snapshots/390b27551542e3b39c89ea96267c3a11f2e00980/transformer",
    )
    transformer = FluxTransformer2DModel.from_pretrained(
        path, torch_dtype=torch.bfloat16, use_safetensors=False
    )
    pipeline = FluxPipeline.from_pretrained(
        ckpt_id,
        revision=ckpt_revision,
        transformer=transformer,
        text_encoder_2=text_encoder_2,
        torch_dtype=torch.bfloat16,
    )
    pipeline.to("cuda")
    pipeline.to(memory_format=torch.channels_last)
    for _ in range(1):
        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


@torch.no_grad()
def infer(
    request: TextToImageRequest, pipeline: Pipeline, generator: Generator
) -> Image:
    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]