Deploy fully customizable image and video pipelines

import modal

MODEL_NAME = "black-forest-labs/FLUX.1-schnell"

image = (

    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04")

    .pip_install("torch", "transformers", "diffusers", ...)

)

volume = modal.Volume.from_name("flux-lora-models")

@app.cls(gpu="H100", image=image, volumes={"/loras": volume})

class FluxWithLoRA:

    @modal.enter()

    def setup(self):

        self.pipeline = FluxPipeline.from_pretrained(MODEL_NAME).to("cuda")

        self.pipeline.load_and_fuse_lora()

    @modal.method()

    def generate_image(self, prompt: str):

        return self.pipeline(prompt).images[0]

flux = FluxWithLoRA()

flux.generate_image.remote("")