Update app.py

app.py

@@ -1,7 +1,7 @@
+#!/usr/bin/env python
 import os
 import random
 import uuid
-import json
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -9,45 +9,55 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-DESCRIPTIONx = """## REALVISXL V5 🤗
-"""
-
 css = '''
-.gradio-container{max-width: 560px !important}
+.gradio-container{max-width: 570px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
 
+DESCRIPTIONXX = """
+    ## REALVISXL V5 + V5 LIGHTNING ⚡
+"""
 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)",
-    "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5, 4K", 
+
+    "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"
 ]
 
-MODEL_ID = os.getenv("MODEL_VAL_PATH", "SG161222/RealVisXL_V4.0") 
+MODEL_OPTIONS = {
+    "REALVISXL V5.0": "SG161222/RealVisXL_V4.0_Lightning",
+    "V5.0 LIGHTNING": "SG161222/RealVisXL_V3.0_Turbo",
+
+}
+
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
-BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  # Allow generating multiple images at once
+BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
-#Load model outside of function
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-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)
-
-# <compile speedup >
-if USE_TORCH_COMPILE:
-    pipe.compile()
-
-# Offloading capacity (RAM)
-if ENABLE_CPU_OFFLOAD:
-    pipe.enable_model_cpu_offload()
+
+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 USE_TORCH_COMPILE:
+        pipe.compile()
+    
+    if ENABLE_CPU_OFFLOAD:
+        pipe.enable_model_cpu_offload()
+    
+    return pipe
+
+# Preload and compile both models
+models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -63,6 +73,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 @spaces.GPU(duration=60, enable_queue=True)
 def generate(
+    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -73,13 +84,15 @@ def generate(
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    num_images: int = 1,  # Number of images to generate
+    num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True),
 ):
+    global models
+    pipe = models[model_choice]
+    
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    #Options 
     options = {
         "prompt": [prompt] * num_images,
         "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
@@ -91,11 +104,9 @@ def generate(
         "output_type": "pil",
     }
 
-    #VRAM usage Lesser
     if use_resolution_binning:
         options["use_resolution_binning"] = True
 
-    #Images potential batches
     images = []
     for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
@@ -106,39 +117,63 @@ def generate(
 
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
-#Main gr.Block
+
+#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(DESCRIPTIONx)  
+    gr.Markdown(DESCRIPTIONXX)
+    with gr.Row():
+        prompt = gr.Text(
+            label="Prompt",
+            show_label=False,
+            max_lines=1,
+            placeholder="Enter your prompt",
+            container=False,
+        )
+        run_button = gr.Button("Run⚡", scale=0)
+    result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
-    with gr.Group():
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(label="Result", columns=1, show_label=False) 
-    with gr.Accordion("Advanced options", open=False, visible=True):
+    with gr.Row():
+        model_choice = gr.Dropdown(
+            label="Model Selection",
+            choices=list(MODEL_OPTIONS.keys()),
+            value="Lightning"
+        )
+
+    with gr.Accordion("Advanced options", open=False, visible==True):
         num_images = gr.Slider(
             label="Number of Images",
             minimum=1,
-            maximum=4,
+            maximum=1,
             step=1,
             value=1,
         )
         with gr.Row():
-            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=5,
-                lines=4,
-                placeholder="Enter a negative prompt",
-                value="bad hands, bad anatomy, ugly, deformed, (face asymmetry, eyes asymmetry, deformed eyes, deformed mouth, open mouth)",
-                visible=True,
-            )
+            with gr.Column(scale=1):
+                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
+                negative_prompt = gr.Text(
+                    label="Negative prompt",
+                    max_lines=5,
+                    lines=4,
+                    placeholder="Enter a negative prompt",
+                    value="(octane render, render, drawing, anime, bad photo, bad photography:1.3), (worst quality, low quality, blurry:1.2), (bad teeth, deformed teeth, deformed lips), (bad anatomy, bad proportions:1.1), (deformed iris, deformed pupils), (deformed eyes, bad eyes), (deformed face, ugly face, bad face), (deformed hands, bad hands, fused fingers), morbid, mutilated, mutation, disfigured",
+                    visible=True,
+                )
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -147,7 +182,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             value=0,
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row(visible=True):
+        with gr.Row():
             width = gr.Slider(
                 label="Width",
                 minimum=512,
@@ -173,7 +208,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             num_inference_steps = gr.Slider(
                 label="Number of inference steps",
                 minimum=1,
-                maximum=25,
+                maximum=40,
                 step=1,
                 value=23,
             )
@@ -190,6 +225,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         outputs=negative_prompt,
         api_name=False,
     )
+    
     gr.on(
         triggers=[
             prompt.submit,
@@ -198,6 +234,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         ],
         fn=generate,
         inputs=[
+            model_choice,
             prompt,
             negative_prompt,
             use_negative_prompt,
@@ -211,7 +248,14 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         ],
         outputs=[result, seed],
         api_name="run",
-    )   
-    
+    )
+    gr.Markdown("🥠Models used in the playground [[REALVISXL V5.0]](https://huggingface.co/SG161222/RealVisXL_V5.0), [[REALVISXL V5.0 LIGHTNING]](https://huggingface.co/SG161222/RealVisXL_V5.0_Lightning) 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 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>.")
+    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())  
+        
 if __name__ == "__main__":
-    demo.queue(max_size=40).launch()
+    demo.queue(max_size=20).launch(show_api=False)