Flux image generator on modal

flux.py

import modal
import asyncio
import io
import time
from typing import List
from pathlib import Path

PATH_ROOT = Path(__file__).parent
HF_HUB_CACHE = "/hf-hub-cache"
MODEL_ID = "black-forest-labs/FLUX.1-schnell"
MODEL_REVISION_ID = "741f7c3ce8b383c54771c7003378a50191e9efe9"
LORA_REPO = "XLabs-AI/flux-RealismLora"

hf_hub_cache_volumes = modal.Volume.from_name(HF_HUB_CACHE[1:], create_if_missing=True)
app = modal.App("flux-app")

container_image = (
    modal.Image.debian_slim(python_version="3.12")
    .add_local_python_source("mymodule", copy=True)
    .pip_install(
        "aiohttp",
        "fastapi[standard]",
        "transformers",
        "diffusers[torch]",
        "hf_transfer",
        "sentencepiece",
        "peft",
    )
    .env(
        {
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
            "HF_HUB_CACHE": HF_HUB_CACHE,
            "DISABLE_TELEMETRY": "YES",
        }
    )
)

with container_image.imports():
    import torch
    from diffusers import FluxPipeline
    from fastapi import Response


@app.cls(
    image=container_image,
    gpu="L40s",  # this minimum enough gpu on modal
    timeout=60,
    container_idle_timeout=60,
    volumes={
        HF_HUB_CACHE: hf_hub_cache_volumes,
    },
)
class Inference:
    @modal.enter()
    def init(self):
        self.pipe = FluxPipeline.from_pretrained(
            MODEL_ID, revision=MODEL_REVISION_ID, torch_dtype=torch.bfloat16
        )
        self.pipe.load_lora_weights(LORA_REPO)
        self.pipe.to("cuda")
        self.generator = torch.Generator("cuda")

    @modal.method()
    def generateImage(
        self,
        prompt="An two astronaut using duck style helmet riding alien predator in mars.",
        width: int = 512,
        height: int = 512,
        seed: int = 423,
        num_inference_steps: int = 4,
        lora_scale: float = 0.5,
    ):
        result = self.pipe(
            prompt=prompt,
            guidance_scale=3.5,
            width=width,
            height=height,
            output_type="pil",
            num_inference_steps=num_inference_steps,
            generator=self.generator.manual_seed(seed),
            joint_attention_kwargs={"scale": lora_scale},
        ).images
        image_output = []
        for image in result:
            with io.BytesIO() as buf:
                image.save(buf, format="PNG")
                image_output.append(buf.getvalue())
        # torch.cuda.empty_cache()
        return image_output

    @modal.web_endpoint()
    def web(
        self,
        prompt: str,
        width: int = 512,
        height: int = 512,
        seed: int = 423,
        num_inference_steps: int = 4,
        lora_scale: float = 0.5,
    ):
        if prompt is None or prompt == "":
            return {"error": "Prompt is required!"}
        result = self.generateImage.local(
            prompt=prompt,
            width=width,
            height=height,
            seed=seed,
            num_inference_steps=num_inference_steps,
        )
        return Response(
            content=result[0],
            media_type="image/png",
        )


@app.local_entrypoint()
def print_ip():
    import random

    inference = Inference()
    start = time.time()
    image: List[any] = inference.generateImage.remote(
        prompt="An indonesian girl, 19 year old, good looking, dark brown dyed hair, dressed pink t-shirt with text 'Flux', blue jeans, walking in metropolis, some people look his bechause its beautyful",
        width=768,
        height=1360,
        num_inference_steps=4,
        seed=random.randint(9999, 99999),
    )
    duration = time.time() + start
    # do save images locally
    print(f"Duration: {duration}s")
    for index, img in enumerate(image):
        with open(f"output_{index}.png", "wb") as f:
            f.write(img)