Update app.py

app.py

@@ -1,12 +1,4 @@
 #!/usr/bin/env python
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# ...
 import os
 import random
 import uuid
@@ -15,7 +7,7 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import FluxPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
 css = '''
 .gradio-container{max-width: 570px !important}
@@ -39,7 +31,7 @@ examples = [
 MODEL_OPTIONS = {
     "Lightning": "SG161222/RealVisXL_V4.0_Lightning",
     "Turbovision": "SG161222/RealVisXL_V3.0_Turbo",
-    "FLUX.1-schnell": "black-forest-labs/FLUX.1-schnell",
+    "FLUX.1-schnell": "black-forest-labs/FLUX.1-schnell"
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -50,11 +42,10 @@ BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 def load_and_prepare_model(model_id):
-    if model_id == "black-forest-labs/FLUX.1-schnell":
-        pipe = DiffusionPipeline.from_pretrained(
+    if "FLUX" in model_id:
+        pipe = FluxPipeline.from_pretrained(
             model_id,
-            torch_dtype=torch.bfloat16,
-            revision="refs/pr/1"
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         ).to(device)
     else:
         pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -266,7 +257,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         outputs=[result, seed],
         api_name="run",
     )
-    gr.Markdown("🥠Models used in the playground [[Lightning]](https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning), [[Turbo]](https://huggingface.co/SG161222/RealVisXL_V3.0_Turbo) for image generation. stable diffusion xl piped (sdxl) model HF. This is the demo space for generating images using the Stable Diffusion XL models, with multi different variants available.")
+    gr.Markdown("🥠Models used in the playground [[Lightning]](https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning), [[Turbo]](https://huggingface.co/SG161222/RealVisXL_V3.0_Turbo), [[FLUX]](https://huggingface.co/black-forest-labs/FLUX.1-schnell) for image generation. stable diffusion xl piped (sdxl) model HF. This is the demo space for generating images using the Stable Diffusion XL models, with multi different variants available.")
     gr.Markdown("🥠This is the demo space for generating images using Stable Diffusion with quality styles, different models and types. Try the sample prompts to generate higher quality images. Try the sample prompts for generating higher quality images.<a href='https://huggingface.co/spaces/prithivMLmods/Top-Prompt-Collection' target='_blank'>Try prompts</a>.")
     gr.Markdown("⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.")
 
@@ -275,4 +266,4 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         predefined_gallery = gr.Gallery(label="Image Gallery", columns=3, show_label=False, value=load_predefined_images())  
         
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(show_api=False)
+    demo.queue(max_size=20).launch(show_api=False)