feat(T2I): add Black Forrest Labs Flux 1 support (livepeer#147)

This commit adds support for the [Black Forrest Labs Flux 1 Schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) model to the T2I pipeline. It is important to note that this model can only run on GPUs with more than 33 GB or VRAM.

runner/app/pipelines/text_to_image.py

@@ -13,7 +13,7 @@
                                  is_turbo_model, split_prompt)
 from diffusers import (AutoPipelineForText2Image, EulerDiscreteScheduler,
                        StableDiffusion3Pipeline, StableDiffusionXLPipeline,
-                       UNet2DConditionModel)
+                       UNet2DConditionModel, FluxPipeline)
 from diffusers.models import AutoencoderKL
 from huggingface_hub import file_download, hf_hub_download
 from safetensors.torch import load_file
@@ -28,6 +28,7 @@ class ModelName(Enum):
     SDXL_LIGHTNING = "ByteDance/SDXL-Lightning"
     SD3_MEDIUM = "stabilityai/stable-diffusion-3-medium-diffusers"
     REALISTIC_VISION_V6 = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
+    FLUX_1_SCHNELL = "black-forest-labs/FLUX.1-schnell"
 
     @classmethod
     def list(cls):
@@ -113,6 +114,12 @@ def __init__(self, model_id: str):
             self.ldm = StableDiffusion3Pipeline.from_pretrained(model_id, **kwargs).to(
                 torch_device
             )
+        elif ModelName.FLUX_1_SCHNELL.value in model_id:
+            # Decrease precision to preven OOM errors.
+            kwargs["torch_dtype"] = torch.bfloat16
+            self.ldm = FluxPipeline.from_pretrained(model_id, **kwargs).to(
+                torch_device
+            )
         else:
             self.ldm = AutoPipelineForText2Image.from_pretrained(model_id, **kwargs).to(
                 torch_device

runner/requirements.txt

@@ -1,4 +1,4 @@
-diffusers==0.29.2
+diffusers==0.30.0
 accelerate==0.30.1
 transformers==4.43.1
 fastapi==0.111.0

runner/test_prompts.py

@@ -0,0 +1,12 @@
+from app.pipelines.utils.utils import split_prompt
+
+if __name__ == "__main__":
+    input_prompt = "A photo of a cat.|"
+    test = split_prompt(input_prompt)
+
+    input_prompt2 = ""
+    test2 = split_prompt(input_prompt2)
+
+    input_pormpt3 = "A photo of a cat.|A photo of a dog.|A photo of a bird."
+    test3 = split_prompt(input_pormpt3)
+