Update app.py

app.py

@@ -10,6 +10,7 @@ import torch
 from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from typing import Tuple
 
+# Load restricted words
 bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
 bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
 default_negative = os.getenv("default_negative", "")
@@ -23,7 +24,7 @@ def check_text(prompt, negative=""):
             return True
     return False
 
-# Quality/Style
+# Quality/Style--------------------------------------------------------------------
 style_list = [
     {
         "name": "3840 x 2160",
@@ -47,7 +48,7 @@ style_list = [
     },
 ]
 
-# Collage styles
+# Collage styles--------------------------------------------------------------------
 collage_style_list = [
     {
         "name": "Hi-Res",
@@ -176,7 +177,7 @@ collage_style_list = [
     },
 ]
 
-# Filters
+# Filters--------------------------------------------------------------------
 filters = {
     "Vivid": {
         "prompt": "extra vivid {prompt}",
@@ -243,10 +244,10 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-# Set dtype based on device: half precision for CUDA, full precision for CPU
+# Set dtype based on device: half for CUDA, float32 for CPU
 dtype = torch.float16 if device.type == "cuda" else torch.float32
 
-# Load the pipeline
+# Load primary model (RealVisXL_V5.0_Lightning)
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "SG161222/RealVisXL_V5.0_Lightning",
@@ -254,27 +255,60 @@ if torch.cuda.is_available():
         use_safetensors=True,
         add_watermarker=False
     ).to(device)
-    # Ensure the text encoder is in half precision to match the rest of the model
+    # Ensure text encoder uses half precision on GPU
     pipe.text_encoder = pipe.text_encoder.half()
     
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
     else:
         pipe.to(device)
-        print("Loaded on Device!")
+        print("Loaded RealVisXL_V5.0_Lightning on Device!")
     
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
+        print("Model RealVisXL_V5.0_Lightning Compiled!")
+    
+    # Load new model (RealVisXL_V4.0)
+    pipe2 = StableDiffusionXLPipeline.from_pretrained(
+        "SG161222/RealVisXL_V4.0",
+        torch_dtype=dtype,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(device)
+    pipe2.text_encoder = pipe2.text_encoder.half()
+    
+    if ENABLE_CPU_OFFLOAD:
+        pipe2.enable_model_cpu_offload()
+    else:
+        pipe2.to(device)
+        print("Loaded RealVisXL_V4.0 on Device!")
+    
+    if USE_TORCH_COMPILE:
+        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
+        print("Model RealVisXL_V4.0 Compiled!")
 else:
-    # On CPU, load with float32
+    # On CPU load both models in float32
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "SG161222/RealVisXL_V5.0_Lightning",
         torch_dtype=dtype,
         use_safetensors=True,
         add_watermarker=False
     ).to(device)
-    print("Running on CPU; model loaded in float32.")
+    pipe2 = StableDiffusionXLPipeline.from_pretrained(
+        "SG161222/RealVisXL_V4.0",
+        torch_dtype=dtype,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(device)
+    print("Model Loaded !!")
+
+# A dictionary to easily choose the model based on selection.
+DEFAULT_MODEL = "RealVisXL_V5.0_Lightning"
+MODEL_CHOICES = [DEFAULT_MODEL, "RealVisXL_V4.0"]
+models = {
+    "RealVisXL_V5.0_Lightning": pipe,
+    "RealVisXL_V4.0": pipe2
+}
 
 def save_image(img, path):
     img.save(path)
@@ -298,6 +332,7 @@ def generate(
     height: int = 1024,
     guidance_scale: float = 3,
     randomize_seed: bool = False,
+    model_choice: str = DEFAULT_MODEL,
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -317,7 +352,7 @@ def generate(
     if not use_negative_prompt:
         negative_prompt = ""
     negative_prompt += default_negative    
-
+    
     grid_sizes = {
         "2x1": (2, 1),
         "1x2": (1, 2),
@@ -343,11 +378,14 @@ def generate(
         "output_type": "pil",
     }
     
-    torch.cuda.empty_cache()  # Clear GPU memory if available
-    images = pipe(**options).images
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    
+    # Choose pipeline based on user selection
+    selected_pipe = models.get(model_choice, pipe)
+    images = selected_pipe(**options).images
 
     grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
-
     for i, img in enumerate(images[:num_images]):
         grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
 
@@ -385,15 +423,20 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            run_button = gr.Button("Generate  as  (1024 x 1024)🍺", scale=0, elem_classes="submit-btn")
-        
+            run_button = gr.Button("Generate as (1024 x 1024)🍺", scale=0, elem_classes="submit-btn")
+            
+            with gr.Row(visible=True):
+                model_selection = gr.Dropdown(
+                    choices=MODEL_CHOICES,
+                    value=DEFAULT_MODEL,
+                    label="Model Selection",
+                )
             with gr.Row(visible=True):
                 grid_size_selection = gr.Dropdown(
                     choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
                     value="1x1",
                     label="Grid Size"
                 )
-
             with gr.Row(visible=True):
                 filter_selection = gr.Dropdown(
                     show_label=True,
@@ -403,7 +446,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     value=DEFAULT_FILTER_NAME,
                     label="Filter Type",
                 )
-
             with gr.Row(visible=True):
                 collage_style_selection = gr.Dropdown(
                     show_label=True,
@@ -413,7 +455,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     value=DEFAULT_COLLAGE_STYLE_NAME,
                     label="Collage Template + Duotone Canvas",
                 )
-
             with gr.Row(visible=True):
                 style_selection = gr.Dropdown(
                     show_label=True,
@@ -423,7 +464,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     value=DEFAULT_STYLE_NAME,
                     label="Quality Style",
                 )
-
             with gr.Accordion("Advanced options", open=False):
                 use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
                 negative_prompt = gr.Text(
@@ -458,7 +498,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     visible=True
                 )
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
                 with gr.Row(visible=True):
                     width = gr.Slider(
                         label="Width",
@@ -474,7 +513,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                         step=64,
                         value=1024,
                     )
-
                 with gr.Row():
                     guidance_scale = gr.Slider(
                         label="Guidance Scale",
@@ -483,10 +521,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                         step=0.1,
                         value=6,
                     )
-
         with gr.Column(scale=2):
             result = gr.Gallery(label="Result", columns=1, show_label=False)
-
             gr.Examples(
                 examples=examples,
                 inputs=prompt,
@@ -494,14 +530,12 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 fn=generate,
                 cache_examples=CACHE_EXAMPLES,
             )
-
         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,
@@ -522,6 +556,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             height,
             guidance_scale,
             randomize_seed,
+            model_selection,
         ],
         outputs=[result, seed],
         api_name="run",