Update pipeline

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from PIL import Image
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
 import gradio as gr
 import spaces
 
@@ -18,7 +18,7 @@ if SUB_MODEL_REPO_ID is None:
 if SUB_MODEL_SUBFOLDER is None:
     raise ValueError("SUB_MODEL_SUBFOLDER is not set")
 
-pipeline = DiffusionPipeline.from_pretrained(
+pipeline = FluxPipeline.from_pretrained(
     MAIN_MODEL_REPO_ID,
     custom_pipeline=SUB_MODEL_REPO_ID,
 ).to(DEVICE)