Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import random
 import uuid
-from typing import Tuple, Dict
+from typing import Tuple
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -9,31 +9,44 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-DESCRIPTIONz= """## SDXL-LoRA-DLC ⚡
-Select a base model, choose a LoRA, and generate images!
+# Description for the Gradio interface
+DESCRIPTIONz = """## SDXL-LoRA-DLC ⚡
 """
 
-# --- Constants ---
+# 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)
+    return seed
+
 MAX_SEED = np.iinfo(np.int32).max
-DEFAULT_STYLE_NAME = "3840 x 2160"
-USE_TORCH_COMPILE = False # Set to True if you want to try torch compile (might be faster but requires compatible hardware/drivers)
-ENABLE_CPU_OFFLOAD = False # Set to True to offload parts of the model to CPU (saves VRAM but slower)
 
-# --- Model Definitions ---
-# Dictionary mapping user-friendly names to Hugging Face model IDs
-pipelines_info = {
+# 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",
-    # Add more SDXL base models here if desired
-    # "Another SDXL Model": "stabilityai/stable-diffusion-xl-base-1.0", # Example
 }
 
-# Dictionary to cache loaded pipelines
-loaded_pipelines: Dict[str, StableDiffusionXLPipeline] = {}
+# Global variables to manage the current pipeline
+current_base_model = None
+current_pipeline = None
 
-# --- LoRA Definitions ---
+# Define LoRA options
 LORA_OPTIONS = {
-    # Name: (HuggingFace Repo ID, Weight Filename, Adapter Name)
     "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"),
@@ -49,88 +62,47 @@ LORA_OPTIONS = {
     "Art Minimalistic (paint/semireal)🎨": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"),
 }
 
-# --- Style Definitions ---
+# Define style options
 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, bad anatomy, worst quality, low quality",
+        "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, bad anatomy, worst quality, low quality",
+        "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, bad anatomy, worst quality, low quality",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
     },
     {
         "name": "Style Zero",
         "prompt": "{prompt}",
-        "negative_prompt": "worst quality, low quality", # Added basic negative 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())
 
-# --- Utility Functions ---
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
+# Function to apply selected style
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    # Get the base style prompt and negative prompt
-    base_p, base_n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-
-    # Combine the base negative prompt with the user's negative prompt
-    # Ensure user's negative prompt is appended correctly
-    if negative and base_n:
-        combined_n = f"{base_n}, {negative}"
-    elif negative:
-        combined_n = negative
-    else:
-        combined_n = base_n
-
-    # Apply the positive prompt template
-    final_p = base_p.replace("{prompt}", positive)
-
-    return final_p, combined_n
-
-def load_predefined_images():
-    # Ensure the assets directory and images exist
-    asset_dir = "assets"
-    image_files = [
-        "1.png", "2.png", "3.png",
-        "4.png", "5.png", "6.png",
-        "7.png", "8.png", "9.png",
-    ]
-    predefined_images = []
-    if os.path.exists(asset_dir):
-        for img_file in image_files:
-            img_path = os.path.join(asset_dir, img_file)
-            if os.path.exists(img_path):
-                predefined_images.append(img_path)
-            else:
-                print(f"Warning: Predefined image not found: {img_path}")
+    if style_name in styles:
+        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     else:
-         print(f"Warning: Asset directory not found: {asset_dir}")
-    # If no images were found, return None or an empty list
-    # to avoid errors in gr.Gallery
-    return predefined_images if predefined_images else None
-
+        p, n = styles[DEFAULT_STYLE_NAME]
+    if not negative:
+        negative = ""
+    return p.replace("{prompt}", positive), n + negative
 
-# --- Core Generation Function ---
+# Generation function with model selection
 @spaces.GPU(duration=180, enable_queue=True)
 def generate(
-    selected_base_model_name: str, # New input for base model selection
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -138,212 +110,110 @@ def generate(
     width: int = 1024,
     height: int = 1024,
     guidance_scale: float = 3,
-    num_inference_steps: int = 4, # Lightning models use fewer steps
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
-    lora_choice: str = "Realism (face/character)👦🏻",
+    lora_model: str = "Realism (face/character)👦🏻",
+    base_model: str = "RealVisXL V5.0 Lightning",
     progress=gr.Progress(track_tqdm=True),
 ):
-    if not torch.cuda.is_available():
-        raise gr.Error("GPU not available. This Space requires a GPU to run.")
+    global current_base_model, current_pipeline
 
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    torch.manual_seed(seed) # Ensure reproducibility if seed is fixed
-
-    # --- Pipeline Loading and Caching ---
-    pipe = None
-    if selected_base_model_name in loaded_pipelines:
-        print(f"Using cached pipeline: {selected_base_model_name}")
-        pipe = loaded_pipelines[selected_base_model_name]
-    else:
-        print(f"Loading pipeline: {selected_base_model_name}")
-        model_id = pipelines_info[selected_base_model_name]
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16" if torch.cuda.is_available() else None # Use fp16 variant if available on GPU
+    # 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
 
-        # Apply optimizations based on flags
-        if ENABLE_CPU_OFFLOAD:
-            print("Enabling CPU Offload")
-            pipe.enable_model_cpu_offload()
-        else:
-             pipe.to("cuda") # Default: move entire pipeline to GPU
-
-        # Configure scheduler (important for Lightning models)
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-
-        # Load ALL LoRAs onto this newly loaded pipeline instance
-        print(f"Loading LoRAs for {selected_base_model_name}...")
-        for lora_name, (model_repo, weight_file, adapter_tag) in LORA_OPTIONS.items():
-             try:
-                 print(f"  Loading LoRA: {lora_name} ({adapter_tag})")
-                 pipe.load_lora_weights(model_repo, weight_name=weight_file, adapter_name=adapter_tag)
-             except Exception as e:
-                 print(f"  Failed to load LoRA {lora_name}: {e}")
-                 # Optionally raise an error or continue without this LoRA
-                 # raise gr.Error(f"Failed to load LoRA {lora_name}. Check repo/file names.")
-
-        if USE_TORCH_COMPILE:
-            print("Attempting to compile UNet (may take time)...")
-            try:
-                pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-                print("UNet compiled successfully.")
-            except Exception as e:
-                print(f"Torch compile failed: {e}. Running without compilation.")
-
-        # Cache the fully loaded and configured pipeline
-        loaded_pipelines[selected_base_model_name] = pipe
-        print(f"Pipeline {selected_base_model_name} loaded and cached.")
-
-    # --- Prompt Styling ---
-    positive_prompt, effective_negative_prompt = apply_style(style_name, prompt, negative_prompt if use_negative_prompt else "")
-
-    # --- LoRA Selection ---
-    if lora_choice not in LORA_OPTIONS:
-         raise gr.Error(f"Selected LoRA '{lora_choice}' not found in options.")
+    # 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 = ""
 
-    _lora_repo, _lora_weight, lora_adapter_name = LORA_OPTIONS[lora_choice]
-    print(f"Activating LoRA: {lora_choice} (Adapter: {lora_adapter_name})")
-    pipe.set_adapters(lora_adapter_name)
-    # Note: LoRA weight/scale is often handled within the pipeline or during loading.
-    # If you need adjustable LoRA scale, you might need `add_weighted_adapter` or similar.
-    # For simplicity here, we assume the default scale is used.
-    # cross_attention_kwargs={"scale": 0.8} # Example if you need to set scale explicitly
+    # Set the LoRA adapter
+    _, _, adapter_name = LORA_OPTIONS[lora_model]
+    current_pipeline.set_adapters(adapter_name)
 
-    # --- Image Generation ---
-    print("Starting image generation...")
-    generator = torch.Generator("cuda").manual_seed(seed)
-    images = pipe(
+    # Generate the image
+    images = current_pipeline(
         prompt=positive_prompt,
         negative_prompt=effective_negative_prompt,
         width=width,
         height=height,
         guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps, # Use steps suitable for Lightning
-        generator=generator,
+        num_inference_steps=20,
         num_images_per_prompt=1,
-        # cross_attention_kwargs=cross_attention_kwargs, # Add if scale needed
+        cross_attention_kwargs={"scale": 0.65},
         output_type="pil",
     ).images
-
     image_paths = [save_image(img) for img in images]
-    print("Image generation complete.")
     return image_paths, seed
 
-# --- Gradio UI ---
+# 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: 860px !important; margin: auto;}
+.gradio-container{max-width: 545px !important}
 h1{text-align:center}
-.gr-prose { text-align: center; }
-#model-select-row { justify-content: center; } /* Center dropdowns */
-/* Make gallery taller */
-#result_gallery .h-\[400px\] {
-    height: 600px !important; /* Adjust height as needed */
+footer {
+    visibility: hidden
 }
-#predefined_gallery .h-\[400px\] {
-    height: 300px !important; /* Adjust height as needed */
-}
-footer { visibility: hidden }
 '''
 
+# Function to load predefined images
+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",
+    ]
+    return predefined_images
+
+# Gradio interface
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTIONz)
-
-    with gr.Row(elem_id="model-select-row"):
-         model_selector = gr.Dropdown(
-             label="Select Base Model",
-             choices=list(pipelines_info.keys()),
-             value=list(pipelines_info.keys())[0], # Default to the first model
-             scale=1
-         )
-         model_choice = gr.Dropdown(
-             label="Select LoRA Style",
-             choices=list(LORA_OPTIONS.keys()),
-             value="Realism (face/character)👦🏻", # Default LoRA
-             scale=1
-         )
-
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=2, # Allow slightly more room for prompt
-                placeholder="Enter your prompt (e.g., 'Astronaut riding a horse')",
+                max_lines=1,
+                placeholder="Enter your prompt with resp. tag!",
                 container=False,
-                scale=5, # Make prompt input wider
             )
-            run_button = gr.Button("Generate", scale=1, variant="primary") # Make button stand out
-
-    # Use Tabs for Main Result and Examples/Gallery
-    with gr.Tabs():
-        with gr.TabItem("Result", id="result_tab"):
-             result = gr.Gallery(
-                 label="Generated Image", elem_id="result_gallery",
-                 columns=1, preview=True, show_label=False, height=600 # Make gallery taller
-             )
-             # Display the seed used for the generated image
-             used_seed = gr.Number(label="Seed Used", interactive=False)
-
-        with gr.TabItem("Examples & Predefined Gallery", id="examples_tab"):
-            gr.Markdown("### Prompt Examples")
-            gr.Examples(
-                examples=[
-                    "cinematic photo, a man sitting on a chair in a dark room, realistic", # Realism example
-                    "pixar style 3d render of a cute cat astronaut exploring mars", # Pixar example
-                    "studio photography, high fashion model wearing a futuristic silver hoodie, dramatic lighting", # Photoshoot/Clothing example
-                    "minimalist vector art illustration of a mountain range at sunset, liquid style", # Minimalist/Liquid example
-                    "pencil sketch drawing of an old wise wizard with a long beard", # Pencil Art example
-                ],
-                inputs=[prompt], # Only update the prompt field from examples
-                outputs=[result, used_seed], # Define outputs for example generation
-                fn=lambda p: generate( # Need a lambda to pass default values for other args
-                    selected_base_model_name=list(pipelines_info.keys())[0], # Use default model for examples
-                    prompt=p,
-                    lora_choice="Realism (face/character)👦🏻", # Use default LoRA for examples
-                    # Add other default args from 'generate' signature if needed
-                    negative_prompt="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-                    use_negative_prompt=True,
-                    seed=0, # Or make examples use random seed?
-                    width=1024,
-                    height=1024,
-                    guidance_scale=3.0,
-                    num_inference_steps=4,
-                    randomize_seed=True, # Randomize seed for examples
-                    style_name=DEFAULT_STYLE_NAME,
-                ),
-                cache_examples=False, # Recalculate examples if needed
-                label="Click an example to generate"
-            )
-            gr.Markdown("### Predefined Image Gallery")
-            predefined_gallery = gr.Gallery(
-                 label="Image Gallery", elem_id="predefined_gallery",
-                 columns=3, show_label=False, value=load_predefined_images(), height=300
-            )
-
-
-    with gr.Accordion("⚙️ Advanced Settings", open=False):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=STYLE_NAMES,
-            value=DEFAULT_STYLE_NAME,
-            label="Image Quality Style",
-        )
-        with gr.Row():
-            use_negative_prompt = gr.Checkbox(label="Use Negative Prompt", value=True)
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+            run_button = gr.Button("Run", scale=0)
+        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
 
+    with gr.Accordion("Advanced options", open=False, visible=False):
+        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
         negative_prompt = gr.Text(
-            label="Negative Prompt",
-            max_lines=2,
-            value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, worst quality, low quality",
-            placeholder="Enter concepts to avoid...",
-            visible=True, # Initially visible, controlled by checkbox change
+            label="Negative prompt",
+            lines=4,
+            max_lines=6,
+            value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
+            placeholder="Enter a negative prompt",
+            visible=True,
         )
         seed = gr.Slider(
             label="Seed",
@@ -351,45 +221,62 @@ with gr.Blocks(css=css) as demo:
             maximum=MAX_SEED,
             step=1,
             value=0,
-            visible=True, # Initially visible, maybe hide if randomize is checked?
-            interactive=True
+            visible=True
         )
-
-        with gr.Row():
+        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
                 minimum=512,
-                maximum=1536, # Adjusted max based on typical SDXL use
-                step=64,
+                maximum=2048,
+                step=8,
                 value=1024,
             )
             height = gr.Slider(
                 label="Height",
                 minimum=512,
-                maximum=1536, # Adjusted max based on typical SDXL use
-                step=64,
+                maximum=2048,
+                step=8,
                 value=1024,
             )
-
         with gr.Row():
             guidance_scale = gr.Slider(
-                label="Guidance Scale (CFG)",
-                minimum=0.0,
-                maximum=10.0, # Lightning models often use low CFG
+                label="Guidance Scale",
+                minimum=0.1,
+                maximum=20.0,
                 step=0.1,
-                value=1.5, # Default low CFG for Lightning
-            )
-            num_inference_steps = gr.Slider(
-                label="Inference Steps",
-                minimum=1,
-                maximum=20, # Lightning models need very few steps
-                step=1,
-                value=4, # Default steps for Lightning
+                value=3.0,
             )
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value=DEFAULT_STYLE_NAME,
+            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"
+        )
+        model_choice = gr.Dropdown(
+            label="LoRA Selection",
+            choices=list(LORA_OPTIONS.keys()),
+            value="Realism (face/character)👦🏻"
+        )
 
-    # --- Event Listeners ---
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=[result, seed],
+        fn=generate,
+        cache_examples=False,
+    )
 
-    # Show/hide negative prompt input based on checkbox
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
@@ -397,54 +284,38 @@ with gr.Blocks(css=css) as demo:
         api_name=False,
     )
 
-    # Show/hide seed slider based on randomize checkbox
-    randomize_seed.change(
-        fn=lambda x: gr.update(interactive=not x), # Make slider non-interactive if randomizing
-        inputs=randomize_seed,
-        outputs=seed,
-        api_name=False,
-    )
-
-    # Main generation trigger
-    inputs_list = [
-        model_selector, # Add model selector
-        prompt,
-        negative_prompt,
-        use_negative_prompt,
-        seed,
-        width,
-        height,
-        guidance_scale,
-        num_inference_steps, # Add steps slider
-        randomize_seed,
-        style_selection,
-        model_choice, # This is the LoRA choice dropdown
-    ]
-    outputs_list = [result, used_seed] # Output gallery and the seed number
-
-    prompt.submit(
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
         fn=generate,
-        inputs=inputs_list,
-        outputs=outputs_list,
-        api_name="run_prompt_submit" # Optional: Define API name
+        inputs=[
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            seed,
+            width,
+            height,
+            guidance_scale,
+            randomize_seed,
+            style_selection,
+            model_choice,
+            base_model_choice,
+        ],
+        outputs=[result, seed],
+        api_name="run",
     )
-    run_button.click(
-        fn=generate,
-        inputs=inputs_list,
-        outputs=outputs_list,
-        api_name="run_button_click" # Optional: Define API name
-    )
-
-# --- Launch ---
-if __name__ == "__main__":
-    if not torch.cuda.is_available():
-        print("Warning: No CUDA GPU detected. Running on CPU will be extremely slow or may fail.")
-        DESCRIPTIONz += "\n<p>⚠️<b>WARNING: No GPU detected. Running on CPU is very slow and may not work reliably.</b> Consider using a GPU instance.</p>"
-        # Optionally disable parts of the UI or exit if CPU is unacceptable
-        # exit()
 
-    # Ensure asset directory exists for predefined images (optional but good practice)
-    if not os.path.exists("assets"):
-        print("Warning: 'assets' directory not found. Predefined images will not load.")
+    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()
+        )
 
-    demo.queue(max_size=20).launch(debug=False) # Set debug=True for more logs if needed
+if __name__ == "__main__":
+    demo.queue(max_size=30).launch()