Update app.py

app.py

@@ -8,10 +8,9 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from diffusers import AuraFlowPipeline
 
 css = '''
-.gradio-container{max-width: 600px !important}
+.gradio-container{max-width: 570px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
@@ -29,7 +28,7 @@ examples = [
 MODEL_OPTIONS = {
     "Lightning": "SG161222/RealVisXL_V4.0_Lightning",
     "Turbovision": "SG161222/RealVisXL_V3.0_Turbo",
-    "AuraFlow": "fal/AuraFlow"
+
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -40,19 +39,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 model_id == "fal/AuraFlow":
-        pipe = AuraFlowPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16 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()
@@ -141,9 +134,6 @@ def load_predefined_images():
     ]
     return predefined_images
 
-
-
-
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     with gr.Row():
         prompt = gr.Text(
@@ -160,7 +150,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         model_choice = gr.Dropdown(
             label="Model Selection",
             choices=list(MODEL_OPTIONS.keys()),
-            value="AuraFlow"
+            value="Lightning"
         )
 
     with gr.Accordion("Advanced options", open=True, visible=False):
@@ -264,7 +254,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
 
     with gr.Column(scale=3):
         gr.Markdown("### Image Gallery")
-        predefined_gallery = gr.Gallery(label="Image Gallery", columns=4, show_label=False, value=load_predefined_images())       
+        predefined_gallery = gr.Gallery(label="Image Gallery", columns=4, 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)