DiffusionPipeline

src/pipeline.py

@@ -4,7 +4,7 @@ from typing import TypeAlias
 
 import torch
 from PIL.Image import Image
-from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, AutoencoderTiny
+from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, AutoencoderTiny, DiffusionPipeline
 from huggingface_hub.constants import HF_HUB_CACHE
 from pipelines.models import TextToImageRequest
 from torch import Generator
@@ -40,7 +40,7 @@ def load_pipeline() -> Pipeline:
             revision=TinyVAE_REV, 
             local_files_only=True,
             torch_dtype=torch.bfloat16)
-    pipeline = FluxPipeline.from_pretrained(
+    pipeline = DiffusionPipeline.from_pretrained(
         CHECKPOINT,
         revision=REVISION,
         transformer=transformer,