Update app.py

app.py

@@ -1,13 +1,3 @@
-#!/usr/bin/env python
-#patch 1.0()
-# 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
-# furnished to do so, subject to the following conditions:
-#
-# ...
 import os
 import random
 import uuid
@@ -19,9 +9,18 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-DESCRIPTIONz= """## SDXL-LoRA-DLC ⚡
+DESCRIPTIONz = """## SDXL-LoRA-DLC ⚡
 """
 
+# Define model options
+MODEL_OPTIONS = {
+    "RealVisXL V4.0 Lightning": "SG161222/RealVisXL_V4.0_Lightning",
+    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
+}
+
+# Dictionary to cache pipelines
+pipelines = {}
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -40,34 +39,36 @@ if not torch.cuda.is_available():
 USE_TORCH_COMPILE = 0
 ENABLE_CPU_OFFLOAD = 0
 
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "SG161222/RealVisXL_V4.0_Lightning",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-
-    LORA_OPTIONS = {
-        "Realism (face/character)👦🏻": ("prithivMLmods/Canopus-Realism-LoRA", "Canopus-Realism-LoRA.safetensors", "rlms"),
-        "Pixar (art/toons)🙀": ("prithivMLmods/Canopus-Pixar-Art", "Canopus-Pixar-Art.safetensors", "pixar"),
-        "Photoshoot (camera/film)📸": ("prithivMLmods/Canopus-Photo-Shoot-Mini-LoRA", "Canopus-Photo-Shoot-Mini-LoRA.safetensors", "photo"),
-        "Clothing (hoodies/pant/shirts)👔": ("prithivMLmods/Canopus-Clothing-Adp-LoRA", "Canopus-Dress-Clothing-LoRA.safetensors", "clth"),
-        "Interior Architecture (house/hotel)🏠": ("prithivMLmods/Canopus-Interior-Architecture-0.1", "Canopus-Interior-Architecture-0.1δ.safetensors", "arch"),
-        "Fashion Product (wearing/usable)👜": ("prithivMLmods/Canopus-Fashion-Product-Dilation", "Canopus-Fashion-Product-Dilation.safetensors", "fashion"),
-        "Minimalistic Image (minimal/detailed)🏞️": ("prithivMLmods/Pegasi-Minimalist-Image-Style", "Pegasi-Minimalist-Image-Style.safetensors", "minimalist"),
-        "Modern Clothing (trend/new)👕": ("prithivMLmods/Canopus-Modern-Clothing-Design", "Canopus-Modern-Clothing-Design.safetensors", "mdrnclth"),
-        "Animaliea (farm/wild)🫎": ("prithivMLmods/Canopus-Animaliea-Artism", "Canopus-Animaliea-Artism.safetensors", "Animaliea"),
-        "Liquid Wallpaper (minimal/illustration)🖼️": ("prithivMLmods/Canopus-Liquid-Wallpaper-Art", "Canopus-Liquid-Wallpaper-Minimalize-LoRA.safetensors", "liquid"),
-        "Canes Cars (realistic/futurecars)🚘": ("prithivMLmods/Canes-Cars-Model-LoRA", "Canes-Cars-Model-LoRA.safetensors", "car"),
-        "Pencil Art (characteristic/creative)✏️": ("prithivMLmods/Canopus-Pencil-Art-LoRA", "Canopus-Pencil-Art-LoRA.safetensors", "Pencil Art"),
-        "Art Minimalistic (paint/semireal)🎨": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"),
-
-    }
-
-    for model_name, weight_name, adapter_name in LORA_OPTIONS.values():
-        pipe.load_lora_weights(model_name, weight_name=weight_name, adapter_name=adapter_name)
-    pipe.to("cuda")
+# Define LoRA options
+LORA_OPTIONS = {
+    "Realism (face/character)👦🏻": ("prithivMLmods/Canopus-Realism-LoRA", "Canopus-Realism-LoRA.safetensors", "rlms"),
+    "Pixar (art/toons)🙀": ("prithivMLmods/Canopus-Pixar-Art", "Canopus-Pixar-Art.safetensors", "pixar"),
+    "Photoshoot (camera/film)📸": ("prithivMLmods/Canopus-Photo-Shoot-Mini-LoRA", "Canopus-Photo-Shoot-Mini-LoRA.safetensors", "photo"),
+    "Clothing (hoodies/pant/shirts)👔": ("prithivMLmods/Canopus-Clothing-Adp-LoRA", "Canopus-Dress-Clothing-LoRA.safetensors", "clth"),
+    "Interior Architecture (house/hotel)🏠": ("prithivMLmods/Canopus-Interior-Architecture-0.1", "Canopus-Interior-Architecture-0.1δ.safetensors", "arch"),
+    "Fashion Product (wearing/usable)👜": ("prithivMLmods/Canopus-Fashion-Product-Dilation", "Canopus-Fashion-Product-Dilation.safetensors", "fashion"),
+    "Minimalistic Image (minimal/detailed)🏞️": ("prithivMLmods/Pegasi-Minimalist-Image-Style", "Pegasi-Minimalist-Image-Style.safetensors", "minimalist"),
+    "Modern Clothing (trend/new)👕": ("prithivMLmods/Canopus-Modern-Clothing-Design", "Canopus-Modern-Clothing-Design.safetensors", "mdrnclth"),
+    "Animaliea (farm/wild)🫎": ("prithivMLmods/Canopus-Animaliea-Artism", "Canopus-Animaliea-Artism.safetensors", "Animaliea"),
+    "Liquid Wallpaper (minimal/illustration)🖼️": ("prithivMLmods/Canopus-Liquid-Wallpaper-Art", "Canopus-Liquid-Wallpaper-Minimalize-LoRA.safetensors", "liquid"),
+    "Canes Cars (realistic/futurecars)🚘": ("prithivMLmods/Canes-Cars-Model-LoRA", "Canes-Cars-Model-LoRA.safetensors", "car"),
+    "Pencil Art (characteristic/creative)✏️": ("prithivMLmods/Canopus-Pencil-Art-LoRA", "Canopus-Pencil-Art-LoRA.safetensors", "Pencil Art"),
+    "Art Minimalistic (paint/semireal)🎨": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"),
+}
+
+# Function to load or retrieve a pipeline
+def get_pipeline(model_name):
+    if model_name not in pipelines:
+        pipelines[model_name] = StableDiffusionXLPipeline.from_pretrained(
+            MODEL_OPTIONS[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipelines[model_name].scheduler = EulerAncestralDiscreteScheduler.from_config(pipelines[model_name].scheduler.config)
+        for lora_model_name, lora_weight_name, lora_adapter_name in LORA_OPTIONS.values():
+            pipelines[model_name].load_lora_weights(lora_model_name, weight_name=lora_weight_name, adapter_name=lora_adapter_name)
+        pipelines[model_name].to("cuda")
+    return pipelines[model_name]
 
 style_list = [
     {
@@ -119,6 +120,7 @@ def generate(
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
     lora_model: str = "Realism (face/character)👦🏻",
+    base_model: str = "RealVisXL V4.0 Lightning",
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
@@ -126,8 +128,9 @@ def generate(
     positive_prompt, effective_negative_prompt = apply_style(style_name, prompt, negative_prompt)
 
     if not use_negative_prompt:
-        effective_negative_prompt = ""  # type: ignore
+        effective_negative_prompt = ""
 
+    pipe = get_pipeline(base_model)
     model_name, weight_name, adapter_name = LORA_OPTIONS[lora_model]
     pipe.set_adapters(adapter_name)
 
@@ -158,10 +161,9 @@ footer {
     visibility: hidden
 }
 '''
+
 def load_predefined_images():
- 
     predefined_images = [
-
         "assets/1.png",
         "assets/2.png",
         "assets/3.png",
@@ -171,7 +173,6 @@ def load_predefined_images():
         "assets/7.png",
         "assets/8.png",
         "assets/9.png",
-
     ]
     return predefined_images
 
@@ -243,6 +244,14 @@ with gr.Blocks(css=css) as demo:
             label="Quality Style",
         )
 
+    # Add base model selection dropdown
+    with gr.Row():
+        base_model_choice = gr.Dropdown(
+            label="Base Model",
+            choices=list(MODEL_OPTIONS.keys()),
+            value="RealVisXL V4.0 Lightning",
+        )
+
     with gr.Row(visible=True):
         model_choice = gr.Dropdown(
             label="LoRA Selection",
@@ -283,15 +292,15 @@ with gr.Blocks(css=css) as demo:
             randomize_seed,
             style_selection,
             model_choice,
+            base_model_choice,
         ],
         outputs=[result, seed],
         api_name="run",
     )
 
-
     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=30).launch()