Update app.py

app.py

@@ -16,14 +16,14 @@ from gradio.events import Events
 from gradio.exceptions import Error
 from gradio_client import utils as client_utils
 from transformers import pipeline
-from diffusers import DiffusionPipeline, FluxPipeline
+from diffusers import DiffusionPipeline
 import torch
 import gradio as gr
 
 # Corrected code with closed parenthesis and explicit token handling
-image_model = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", 
-    torch_dtype=torch.bfloat16, 
+image_model = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
     use_auth_token=os.getenv("HUGGINGFACE_TOKEN")
 )
 image_model.enable_model_cpu_offload()