Update app.py

app.py

@@ -9,17 +9,14 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-# Description for the Gradio interface
-DESCRIPTIONz = """## SDXL-LoRA-DLC ⚡
+DESCRIPTIONz= """## SDXL-LoRA-DLC ⚡
 """
 
-# Function to save generated images
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
-# Function to handle seed randomization
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -27,42 +24,41 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 MAX_SEED = np.iinfo(np.int32).max
 
-# Warning if running on CPU
 if not torch.cuda.is_available():
     DESCRIPTIONz += "\n<p>⚠️Running on CPU, This may not work on CPU. If it runs for an extended time or if you encounter errors, try running it on a GPU by duplicating the space using @spaces.GPU(). +import spaces.📍</p>"
 
-# Configuration flags (unchanged)
 USE_TORCH_COMPILE = 0
 ENABLE_CPU_OFFLOAD = 0
 
-# Define available base models
-base_models = {
-    "RealVisXL V4.0 Lightning": "SG161222/RealVisXL_V4.0_Lightning",
-    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
-}
+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)
 
-# Global variables to manage the current pipeline
-current_base_model = None
-current_pipeline = None
-
-# 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"),
-}
+    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 style options
 style_list = [
     {
         "name": "3840 x 2160",
@@ -87,20 +83,20 @@ style_list = [
 ]
 
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+
 DEFAULT_STYLE_NAME = "3840 x 2160"
 STYLE_NAMES = list(styles.keys())
 
-# Function to apply selected style
 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
 
-# Generation function with model selection
 @spaces.GPU(duration=180, enable_queue=True)
 def generate(
     prompt: str,
@@ -113,39 +109,19 @@ def generate(
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
     lora_model: str = "Realism (face/character)👦🏻",
-    base_model: str = "RealVisXL V5.0 Lightning",
     progress=gr.Progress(track_tqdm=True),
 ):
-    global current_base_model, current_pipeline
-
-    # Load the pipeline if the base model has changed
-    if base_model != current_base_model:
-        model_id = base_models[base_model]
-        current_pipeline = StableDiffusionXLPipeline.from_pretrained(
-            model_id, torch_dtype=torch.float16, use_safetensors=True
-        )
-        current_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-            current_pipeline.scheduler.config
-        )
-        for lora_display_name, (lora_model, lora_weight, adapter_name) in LORA_OPTIONS.items():
-            current_pipeline.load_lora_weights(
-                lora_model, weight_name=lora_weight, adapter_name=adapter_name
-            )
-        current_pipeline.to("cuda")
-        current_base_model = base_model
-
-    # Handle seed and prompts
     seed = int(randomize_seed_fn(seed, randomize_seed))
+
     positive_prompt, effective_negative_prompt = apply_style(style_name, prompt, negative_prompt)
+
     if not use_negative_prompt:
-        effective_negative_prompt = ""
+        effective_negative_prompt = ""  # type: ignore
 
-    # Set the LoRA adapter
-    _, _, adapter_name = LORA_OPTIONS[lora_model]
-    current_pipeline.set_adapters(adapter_name)
+    model_name, weight_name, adapter_name = LORA_OPTIONS[lora_model]
+    pipe.set_adapters(adapter_name)
 
-    # Generate the image
-    images = current_pipeline(
+    images = pipe(
         prompt=positive_prompt,
         negative_prompt=effective_negative_prompt,
         width=width,
@@ -159,14 +135,12 @@ def generate(
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-# Example prompts
 examples = [
     "Realism: Man in the style of dark beige and brown, uhd image, youthful protagonists, nonrepresentational ",
     "Pixar: A young man with light brown wavy hair and light brown eyes sitting in an armchair and looking directly at the camera, pixar style, disney pixar, office background, ultra detailed, 1 man",
     "Hoodie: Front view, capture a urban style, Superman Hoodie, technical materials, fabric small point label on text Blue theory, the design is minimal, with a raised collar, fabric is a Light yellow, low angle to capture the Hoodies form and detailing, f/5.6 to focus on the hoodies craftsmanship, solid grey background, studio light setting, with batman logo in the chest region of the t-shirt",
 ]
 
-# CSS styling
 css = '''
 .gradio-container{max-width: 545px !important}
 h1{text-align:center}
@@ -174,10 +148,10 @@ footer {
     visibility: hidden
 }
 '''
-
-# Function to load predefined images
 def load_predefined_images():
+ 
     predefined_images = [
+
         "assets/1.png",
         "assets/2.png",
         "assets/3.png",
@@ -187,10 +161,10 @@ def load_predefined_images():
         "assets/7.png",
         "assets/8.png",
         "assets/9.png",
+
     ]
     return predefined_images
 
-# Gradio interface
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTIONz)
     with gr.Group():
@@ -224,6 +198,7 @@ with gr.Blocks(css=css) as demo:
             visible=True
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
         with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
@@ -239,6 +214,7 @@ with gr.Blocks(css=css) as demo:
                 step=8,
                 value=1024,
             )
+
         with gr.Row():
             guidance_scale = gr.Slider(
                 label="Guidance Scale",
@@ -247,6 +223,7 @@ with gr.Blocks(css=css) as demo:
                 step=0.1,
                 value=3.0,
             )
+
         style_selection = gr.Radio(
             show_label=True,
             container=True,
@@ -256,13 +233,7 @@ with gr.Blocks(css=css) as demo:
             label="Quality Style",
         )
 
-    # Add base model and LoRA selection dropdowns
-    with gr.Row():
-        base_model_choice = gr.Dropdown(
-            label="Base Model",
-            choices=list(base_models.keys()),
-            value="RealVisXL V5.0 Lightning"
-        )
+    with gr.Row(visible=True):
         model_choice = gr.Dropdown(
             label="LoRA Selection",
             choices=list(LORA_OPTIONS.keys()),
@@ -302,20 +273,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()
-        )
-
+        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()