Update app.py

app.py

@@ -1,4 +1,12 @@
 #!/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
@@ -7,7 +15,7 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
 css = '''
 .gradio-container{max-width: 570px !important}
@@ -18,9 +26,7 @@ footer {
 '''
 
 DESCRIPTIONXX = """
-
     ## TEXT 2 IMAGE🥠
-
 """
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -33,7 +39,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"))
@@ -44,13 +50,20 @@ 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):
-    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 model_id == "black-forest-labs/FLUX.1-schnell":
+        pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            revision="refs/pr/1"
+        ).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)
     
     if USE_TORCH_COMPILE:
         pipe.compile()
@@ -60,7 +73,7 @@ def load_and_prepare_model(model_id):
     
     return pipe
 
-# Preload and compile both models
+# Preload and compile all models
 models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
 
 MAX_SEED = np.iinfo(np.int32).max