Update app.py

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import AuraFlowPipeline
 
 css = '''
 .gradio-container{max-width: 570px !important}
@@ -28,7 +29,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"))
@@ -39,13 +40,19 @@ 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()
@@ -117,31 +124,27 @@ def generate(
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-#def load_predefined_images():
-  #  predefined_images = [
- #       "assets/1.png",
-    #    "assets/2.png",
-    #    "assets/3.png",
-    #    "assets/4.png",
-    #    "assets/5.png",
-   #     "assets/6.png",
-    #    "assets/7.png",
-     #   "assets/8.png",
-    #    "assets/9.png",
-     #   "assets/10.png",
-   #     "assets/11.png",
- #       "assets/12.png",
- #   ]
-  #  return predefined_images
+def load_predefined_images():
+    predefined_images = [
+        "assets/1.png",
+        "assets/2.png",
+        "assets/3.png",
+        "assets/4.png",
+        "assets/5.png",
+        "assets/6.png",
+        "assets/7.png",
+        "assets/8.png",
+        "assets/9.png",
+        "assets/10.png",
+        "assets/11.png",
+        "assets/12.png",
+    ]
+    return predefined_images
+
+
 
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown(
-            f"""
-        # Text🥠Image
-        Models used in the playground [[Lightning]](https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning), [[Realvision]](https://huggingface.co/) ,[[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. ⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.
-        """
-        )
     with gr.Row():
         prompt = gr.Text(
             label="Prompt",
@@ -157,7 +160,7 @@ with gr.Blocks(css=css) as demo:
         model_choice = gr.Dropdown(
             label="Model Selection",
             choices=list(MODEL_OPTIONS.keys()),
-            value="Lightning"
+            value="AuraFlow"
         )
 
     with gr.Accordion("Advanced options", open=True, visible=False):
@@ -254,8 +257,14 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed],
         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())       
+    gr.Markdown("🥠Models used in the playground [[Lightning]](https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning), [[AuraFlow]](https://huggingface.co/fal/AuraFlow) ,[[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. ⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.")
+
+    gr.Markdown("🥠This is the demo space for generating images using Stable Diffusion with grids, filters, templates, quality styles, 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.")
+
+    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=20).launch(show_api=False)