Update app.py

app.py

@@ -8,10 +8,9 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from fal import AuraFlowPipeline
 
 DESCRIPTIONx = """
-
-
 """
 
 css = '''
@@ -22,17 +21,10 @@ footer {
 }
 '''
 
-#examples = [
- #   "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
-  #  "Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic oil --ar 2:3 --q 2 --s 750 --v 5  --ar 2:3 --q 2 --s 750 --v 5",
-  #  "Illustration of A starry night camp in the mountains. Low-angle view, Minimal background, Geometric shapes theme, Pottery, Split-complementary colors, Bicolored light, UHD",
-  #  "Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5",
-   # "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 "
-#]
-
 MODEL_OPTIONS = {
     "Lightning": "SG161222/RealVisXL_V4.0_Lightning",
     "Realvision": "SG161222/RealVisXL_V4.0",
+    "AuraFlow": "fal/AuraFlow",
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -43,23 +35,29 @@ 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 == "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)
+
     if USE_TORCH_COMPILE:
         pipe.compile()
-    
+
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
-    
+
     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
@@ -92,7 +90,7 @@ def generate(
 ):
     global models
     pipe = models[model_choice]
-    
+
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -139,7 +137,7 @@ def load_predefined_images():
     return predefined_images
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTIONx)    
+    gr.Markdown(DESCRIPTIONx)
     with gr.Row():
         prompt = gr.Text(
             label="Prompt",
@@ -150,7 +148,7 @@ with gr.Blocks(css=css) as demo:
             container=False,
         )
         run_button = gr.Button("Run⚡", scale=0)
-    result = gr.Gallery(label="Result", columns=1, show_label=False) 
+    result = gr.Gallery(label="Result", columns=1, show_label=False)
 
     with gr.Row():
         model_choice = gr.Dropdown(
@@ -217,19 +215,13 @@ with gr.Blocks(css=css) as demo:
                 value=20,
             )
 
-   # gr.Examples(
-    #    examples=examples,
-    #    inputs=prompt,
-     #   cache_examples=False
-    #)
-
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
         outputs=negative_prompt,
         api_name=False,
     )
-    
+
     gr.on(
         triggers=[
             prompt.submit,
@@ -254,9 +246,5 @@ with gr.Blocks(css=css) as demo:
         api_name="run",
     )
 
- #   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())
-           
 if __name__ == "__main__":
     demo.queue(max_size=40).launch(show_api=False)