flux pipe

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import numpy as np
 import random
 #import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, FluxPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -13,7 +13,7 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.float16)
+pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe.unet.load_attn_procs("enhanceaiteam/Flux-uncensored-V2")
 pipe = pipe.to(device)