Update app.py

app.py

@@ -1,4 +1,10 @@
 #!/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
+
 import os
 import random
 import uuid
@@ -8,6 +14,7 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from typing import Tuple
 
 css = '''
 .gradio-container{max-width: 570px !important}
@@ -18,18 +25,22 @@ footer {
 '''
 
 DESCRIPTIONXX = """
-    ## TEXT 2 IMAGE🥠
+    ## MidJourney XL🏜️
 """
+
 examples = [
 
-    "Illustration of A starry night camp in the mountains, 4k, cinematic --ar 85:128 --v 6.0 --style raw",
-    "A delicious ceviche cheesecake slice, 4k, octane render, ray tracing, Ultra-High-Definition"
+    "Many apples splashed with drops of water within a fancy bowl 4k, hdr  --v 6.0 --style raw",
+    "A profile photo of a dog, brown background, shot on Leica M6 --ar 128:85 --v 6.0 --style raw",
 ]
 
 MODEL_OPTIONS = {
-    "Lightning": "SG161222/RealVisXL_V4.0_Lightning",
-    "Turbovision": "SG161222/RealVisXL_V3.0_Turbo",
-
+    
+    "REALVISXL V5.0": "SG161222/RealVisXL_V5.0",
+    "LIGHTNING V5.0": "SG161222/RealVisXL_V5.0_Lightning",
+    "REALVISXL V4.0": "SG161222/RealVisXL_V4.0",
+    "LIGHTNING V4.0": "SG161222/RealVisXL_V4.0_Lightning",
+    "REALVISXL V3.0 TURBO": "SG161222/RealVisXL_V3.0_Turbo",
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -39,6 +50,43 @@ BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+style_list = [
+    {
+        "name": "3840 x 2160",
+        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "2560 x 1440",
+        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "HD+",
+        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "Style Zero",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
+]
+
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+DEFAULT_STYLE_NAME = "3840 x 2160"
+STYLE_NAMES = list(styles.keys())
+
+def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
+    if style_name in styles:
+        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    else:
+        p, n = styles[DEFAULT_STYLE_NAME]
+
+    if not negative:
+        negative = ""
+    return p.replace("{prompt}", positive), n + negative
+
 def load_and_prepare_model(model_id):
     pipe = StableDiffusionXLPipeline.from_pretrained(
         model_id,
@@ -77,6 +125,7 @@ def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
+    style_selection: str = DEFAULT_STYLE_NAME,
     seed: int = 1,
     width: int = 1024,
     height: int = 1024,
@@ -93,6 +142,8 @@ def generate(
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
+    prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
+
     options = {
         "prompt": [prompt] * num_images,
         "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
@@ -118,22 +169,6 @@ 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
 
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown(DESCRIPTIONXX)
@@ -145,21 +180,29 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             placeholder="Enter your prompt",
             container=False,
         )
-        run_button = gr.Button("Run⚡", scale=0)
+        run_button = gr.Button("Run", scale=0)
     result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
     with gr.Row():
         model_choice = gr.Dropdown(
-            label="Model Selection",
+            label="Model Selection🔻",
             choices=list(MODEL_OPTIONS.keys()),
-            value="Lightning"
+            value="LIGHTNING V5.0"
         )
 
-    with gr.Accordion("Advanced options", open=True, visible=False):
+    with gr.Accordion("Advanced options", open=False, visible=False):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value=DEFAULT_STYLE_NAME,
+            label="Quality Style",
+        )
         num_images = gr.Slider(
             label="Number of Images",
             minimum=1,
-            maximum=1,
+            maximum=5,
             step=1,
             value=1,
         )
@@ -208,9 +251,9 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             num_inference_steps = gr.Slider(
                 label="Number of inference steps",
                 minimum=1,
-                maximum=35,
+                maximum=60,
                 step=1,
-                value=20,
+                value=28,
             )
 
     gr.Examples(
@@ -238,24 +281,41 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             prompt,
             negative_prompt,
             use_negative_prompt,
+            style_selection,
             seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
-            randomize_seed,
-            num_images
+            randomize_seed, 
+            num_images, 
         ],
         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) 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.")
-
-    with gr.Column(scale=3):
-        gr.Markdown("### Image Gallery")
-        predefined_gallery = gr.Gallery(label="Image Gallery", columns=3, show_label=False, value=load_predefined_images())  
-        
+
+    gr.Markdown(
+    """
+    <div style="text-align: justify;">
+    ⚡Models used in the playground: <a href="https://huggingface.co/SG161222/RealVisXL_V5.0">[REALVISXL V5.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V5.0_Lightning">[LIGHTNING V5.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V4.0">[REALVISXL V4.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning">[LIGHTNING V4.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V3.0_Turbo">[REALVISXL V3.0 TURBO]</a>
+    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 multiple different variants available.
+    </div>
+    """
+    )
+
+    gr.Markdown(
+    """
+    <div style="text-align: justify;">
+    ⚡This is the demo space for generating images using Stable Diffusion XL 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>.
+    </div>
+    """)
+
+    gr.Markdown(
+    """
+    <div style="text-align: justify;">
+    ⚠️ Users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.
+    </div>
+    """) 
+    
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(show_api=False)
+    demo.queue(max_size=50).launch(show_api=False)