Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import spaces
 import torch
 import random
-from diffusers import DiffusionPipeline, AutoPipelineForText2Image
+from diffusers import DiffusionPipeline, AutoPipelineForText2Image, FluxPipeline
 import os
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -19,7 +19,7 @@ def generate(prompt, negative_prompt, model=image_model):
         negative_prompt = "ugly, deformed, disfigured, poor quality, low resolution"
     if model == 'enhanceaiteam/Flux-uncensored-v2':
         # Load the base model
-        pipe = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=hf_token).to('cuda')
+        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=hf_token).to('cuda')
         
         # Load the uncensored LoRA weights
         pipe.load_lora_weights('enhanceaiteam/Flux-uncensored-v2', weight_name='lora.safetensors')