Update app.py

app.py

@@ -7,7 +7,7 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import FluxPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
 css = '''
 .gradio-container{max-width: 570px !important}
@@ -31,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"
+
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -42,19 +42,13 @@ 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 "FLUX" in model_id:
-        pipe = FluxPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        ).to(device)
-    else:
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            use_safetensors=True,
-            add_watermarker=False,
-        ).to(device)
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(device)
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     
     if USE_TORCH_COMPILE:
         pipe.compile()
@@ -64,7 +58,7 @@ def load_and_prepare_model(model_id):
     
     return pipe
 
-# Preload and compile all models
+# Preload and compile both models
 models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -257,7 +251,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), [[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("🥠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("🥠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.")
 
@@ -266,4 +260,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)