feat: multi version

app.py

@@ -13,7 +13,7 @@ from datetime import datetime
 from diffusers.models import AutoencoderKL
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
 from config import (
-    MODEL,
+    MODELS,
     MIN_IMAGE_SIZE,
     MAX_IMAGE_SIZE,
     USE_TORCH_COMPILE,
@@ -90,6 +90,7 @@ def generate(
     guidance_scale: float = 6.0,
     num_inference_steps: int = 25,
     sampler: str = "Euler a",
+    model_name: str = "v14",
     aspect_ratio_selector: str = DEFAULT_ASPECT_RATIO,
     use_upscaler: bool = False,
     upscaler_strength: float = 0.55,
@@ -133,6 +134,7 @@ def generate(
         width, height = utils.preprocess_image_dimensions(width, height)
 
         # Set up pipeline
+        pipe = pipes[model_name]
         backup_scheduler = pipe.scheduler
         pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
 
@@ -148,7 +150,7 @@ def generate(
             "num_inference_steps": num_inference_steps,
             "seed": seed,
             "sampler": sampler,
-            "Model": "WAI NSFW illustrious SDXL v14",
+            "Model": "WAI NSFW illustrious SDXL " + model_name,
             "Model hash": "BDB59BAC77D94AE7A55FF893170F9554C3F349E48A1B73C0C17C0B7C6F4D41A2",
         }
 
@@ -238,17 +240,26 @@ def generate(
 
 # Model initialization
 if torch.cuda.is_available():
+    pipes = {}
+
     try:
         logger.info("Loading VAE and pipeline...")
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix",
             torch_dtype=torch.float16,
         )
-        pipe = utils.load_pipeline(MODEL, device, vae=vae)
+
+        for model in MODELS:
+            pipes[model['name']] = utils.load_pipeline(model['path'], device, vae=vae)
+            logger.info(f"Pipeline for {model} loaded successfully on GPU!")
         logger.info("Pipeline loaded successfully on GPU!")
     except Exception as e:
         logger.error(f"Error loading VAE, falling back to default: {e}")
-        pipe = utils.load_pipeline(MODEL, device)
+
+        for model in MODELS:
+            if model['name'] not in pipes:
+                pipes[model['name']] = utils.load_pipeline(model['path'], device)
+                logger.info(f"Pipeline for {model} loaded successfully on GPU!")
 else:
     logger.warning("CUDA not available, running on CPU")
     # check if os.environ keys have VSCODE, if not, load the model on CPU
@@ -260,7 +271,7 @@ else:
 
     if not skip:
         logger.info("Loading pipeline on CPU...")
-        pipe = utils.load_pipeline(MODEL, torch.device("cpu"))
+        pipe = utils.load_pipeline(MODELS[0]['name'], torch.device("cpu"))
         logger.info("Pipeline loaded successfully on CPU!")
 
 
@@ -296,6 +307,13 @@ with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme_5") as demo:
                     value=True,
                     info="Add quality-enhancing tags to your prompt automatically.",
                 )
+                model_name = gr.Radio(
+                        label="Model",
+                        choices=[model['name'] for model in MODELS],
+                        value=MODELS[0]['name'],
+                        container=True,
+                        info="Select the model to use for image generation.",
+                    )
             with gr.Accordion(label="More Settings", open=False):
                 with gr.Group():
                     aspect_ratio_selector = gr.Radio(
@@ -454,6 +472,7 @@ with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme_5") as demo:
             guidance_scale,
             num_inference_steps,
             sampler,
+            model_name,
             aspect_ratio_selector,
             use_upscaler,
             upscaler_strength,

config.py

@@ -26,7 +26,7 @@ config = load_config()
 css = load_css()
 
 # Export variables for backward compatibility
-MODEL = os.getenv("MODEL", config['model']['path'])
+MODELS = os.getenv("MODELS", config['models'])
 MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", config['model']['min_image_size']))
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", config['model']['max_image_size']))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", str(config['model']['use_torch_compile'])).lower() == "true"

config.toml

@@ -1,11 +1,30 @@
 [model]
-path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v14.safetensors"
 min_image_size = 512
 max_image_size = 2880
 use_torch_compile = false
 enable_cpu_offload = false
 output_dir = "./outputs"
 
+[[models]]
+path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v14.safetensors"
+name = "v14"
+
+[[models]]
+path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v13.safetensors"
+name = "v13"
+
+[[models]]
+path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v12.safetensors"
+name = "v12"
+
+[[models]]
+path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v11.safetensors"
+name = "v11"
+
+[[models]]
+path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v10.safetensors"
+name = "v10"
+
 [prompts]
 quality_tags = "masterpiece,best quality,amazing quality, {prompt}"
 default_negative = "bad quality,worst quality,worst detail,sketch,censor"