migrating to zero gpu

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -4,13 +4,11 @@ emoji: 🌍
 colorFrom: gray
 colorTo: purple
 sdk: gradio
-sdk_version: 4.2.0
+sdk_version: 4.20.0
 app_file: app.py
 license: mit
 pinned: false
 suggested_hardware: a10g-small
-duplicated_from: hysts/SD-XL
-hf_oauth: true
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,244 +1,71 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
 import os
-import random
 import gc
-import toml
 import gradio as gr
 import numpy as np
-import utils
 import torch
 import json
-import PIL.Image
-import base64
-import safetensors
-from io import BytesIO
-from typing import Tuple
+import spaces
+import config
+import utils
+import logging
+from PIL import Image, PngImagePlugin
 from datetime import datetime
-from PIL import PngImagePlugin
-import gradio_user_history as gr_user_history
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from lora_diffusers import LoRANetwork, create_network_from_weights
 from diffusers.models import AutoencoderKL
-from diffusers import (
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
-    DPMSolverMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2DiscreteScheduler,
-    EulerDiscreteScheduler,
-    EulerAncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-    LMSDiscreteScheduler,
-    DDIMScheduler,
-    DEISMultistepScheduler,
-    UniPCMultistepScheduler,
-)
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 DESCRIPTION = "Animagine XL 3.0"
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
 IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
-MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN = os.getenv("HF_TOKEN")
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
 MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
 
-MODEL = os.getenv("MODEL", "https://huggingface.co/cagliostrolab/animagine-xl-3.0/blob/main/animagine-xl-3.0.safetensors")
+MODEL = os.getenv(
+    "MODEL",
+    "https://huggingface.co/cagliostrolab/animagine-xl-3.0/blob/main/animagine-xl-3.0.safetensors",
+)
 
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-if torch.cuda.is_available():
+
+def load_pipeline(model_name):
     vae = AutoencoderKL.from_pretrained(
         "madebyollin/sdxl-vae-fp16-fix",
         torch_dtype=torch.float16,
     )
-    pipeline = StableDiffusionXLPipeline.from_single_file if MODEL.endswith(".safetensors") else StableDiffusionXLPipeline.from_pretrained
-    
+    pipeline = (
+        StableDiffusionXLPipeline.from_single_file
+        if MODEL.endswith(".safetensors")
+        else StableDiffusionXLPipeline.from_pretrained
+    )
+
     pipe = pipeline(
-        MODEL,
+        model_name,
         vae=vae,
         torch_dtype=torch.float16,
         custom_pipeline="lpw_stable_diffusion_xl",
         use_safetensors=True,
+        add_watermarker=False,
         use_auth_token=HF_TOKEN,
         variant="fp16",
     )
 
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-    else:
-        pipe.to(device)
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-else:
-    pipe = None
-
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-
-def seed_everything(seed):
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    generator = torch.Generator()
-    generator.manual_seed(seed)
-    return generator
-
-
-def get_image_path(base_path: str):
-    extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif"]
-    for ext in extensions:
-        image_path = base_path + ext
-        if os.path.exists(image_path):
-            return image_path
-    return None
-
-
-def update_selection(selected_state: gr.SelectData):
-    lora_repo = sdxl_loras[selected_state.index]["repo"]
-    lora_weight = sdxl_loras[selected_state.index]["multiplier"]
-    updated_selected_info = f"{lora_repo}"
-
-    return (
-        updated_selected_info,
-        selected_state,
-        lora_weight,
-    )
-
-
-def parse_aspect_ratio(aspect_ratio):
-    if aspect_ratio == "Custom":
-        return None, None
-    width, height = aspect_ratio.split(" x ")
-    return int(width), int(height)
-
-
-def aspect_ratio_handler(aspect_ratio, custom_width, custom_height):
-    if aspect_ratio == "Custom":
-        return custom_width, custom_height
-    else:
-        width, height = parse_aspect_ratio(aspect_ratio)
-        return width, height
-
-
-def create_network(text_encoders, unet, state_dict, multiplier, device):
-    network = create_network_from_weights(
-        text_encoders,
-        unet,
-        state_dict,
-        multiplier,
-    )
-    network.load_state_dict(state_dict)
-    network.to(device, dtype=unet.dtype)
-    network.apply_to(multiplier=multiplier)
-
-    return network
-
-
-def get_scheduler(scheduler_config, name):
-    scheduler_map = {
-        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
-            scheduler_config, use_karras_sigmas=True
-        ),
-        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
-            scheduler_config, use_karras_sigmas=True
-        ),
-        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
-            scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
-        ),
-        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
-        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
-            scheduler_config
-        ),
-        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
-    }
-    return scheduler_map.get(name, lambda: None)()
-
-
-def free_memory():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
-def preprocess_prompt(
-    style_dict,
-    style_name: str,
-    positive: str,
-    negative: str = "",
-    add_style: bool = True,
-) -> Tuple[str, str]:
-    p, n = style_dict.get(style_name, style_dict["(None)"])
-
-    if add_style and positive.strip():
-        formatted_positive = p.format(prompt=positive)
-    else:
-        formatted_positive = positive
-
-    combined_negative = n + negative
-    return formatted_positive, combined_negative
-
-
-def common_upscale(samples, width, height, upscale_method):
-    return torch.nn.functional.interpolate(
-        samples, size=(height, width), mode=upscale_method
-    )
-
-
-def upscale(samples, upscale_method, scale_by):
-    width = round(samples.shape[3] * scale_by)
-    height = round(samples.shape[2] * scale_by)
-    s = common_upscale(samples, width, height, upscale_method)
-    return s
-
+    pipe.to(device)
+    return pipe
 
-def load_and_convert_thumbnail(model_path: str):
-    with safetensors.safe_open(model_path, framework="pt") as f:
-        metadata = f.metadata()
-    if "modelspec.thumbnail" in metadata:
-        base64_data = metadata["modelspec.thumbnail"]
-        prefix, encoded = base64_data.split(",", 1)
-        image_data = base64.b64decode(encoded)
-        image = PIL.Image.open(BytesIO(image_data))
-        return image
-    return None
-
-def load_wildcard_files(wildcard_dir):
-    wildcard_files = {}
-    for file in os.listdir(wildcard_dir):
-        if file.endswith(".txt"):
-            key = f"__{file.split('.')[0]}__"  # Create a key like __character__
-            wildcard_files[key] = os.path.join(wildcard_dir, file)
-    return wildcard_files
-
-def get_random_line_from_file(file_path):
-    with open(file_path, 'r') as file:
-        lines = file.readlines()
-        if not lines:
-            return ""
-        return random.choice(lines).strip()
-
-def add_wildcard(prompt, wildcard_files):
-    for key, file_path in wildcard_files.items():
-        if key in prompt:
-            wildcard_line = get_random_line_from_file(file_path)
-            prompt = prompt.replace(key, wildcard_line)
-    return prompt
 
+@spaces.GPU
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -247,90 +74,40 @@ def generate(
     custom_height: int = 1024,
     guidance_scale: float = 7.0,
     num_inference_steps: int = 28,
-    use_lora: bool = False,
-    lora_weight: float = 1.0,
-    selected_state: str = "",
     sampler: str = "Euler a",
     aspect_ratio_selector: str = "896 x 1152",
     style_selector: str = "(None)",
     quality_selector: str = "Standard",
     use_upscaler: bool = False,
-    upscaler_strength: float = 0.5,
+    upscaler_strength: float = 0.55,
     upscale_by: float = 1.5,
     add_quality_tags: bool = True,
-    profile: gr.OAuthProfile | None = None,
     progress=gr.Progress(track_tqdm=True),
-) -> PIL.Image.Image:
-    generator = seed_everything(seed)
+) -> Image:
+    generator = utils.seed_everything(seed)
 
-    network = None
-    network_state = {"current_lora": None, "multiplier": None}
-
-    width, height = aspect_ratio_handler(
+    width, height = utils.aspect_ratio_handler(
         aspect_ratio_selector,
         custom_width,
         custom_height,
     )
 
-    prompt = add_wildcard(prompt, wildcard_files)
+    prompt = utils.add_wildcard(prompt, wildcard_files)
 
-    
-    prompt, negative_prompt = preprocess_prompt(
+    prompt, negative_prompt = utils.preprocess_prompt(
         quality_prompt, quality_selector, prompt, negative_prompt, add_quality_tags
     )
-    prompt, negative_prompt = preprocess_prompt(
+    prompt, negative_prompt = utils.preprocess_prompt(
         styles, style_selector, prompt, negative_prompt
     )
 
-    if width % 8 != 0:
-        width = width - (width % 8)
-    if height % 8 != 0:
-        height = height - (height % 8)
-        
-    if use_lora:
-        if not selected_state:
-            raise Exception("You must Select a LoRA")
-        repo_name = sdxl_loras[selected_state.index]["repo"]
-        full_path_lora = saved_names[selected_state.index]
-        weight_name = sdxl_loras[selected_state.index]["weights"]
-
-        lora_sd = load_file(full_path_lora)
-        text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
-
-        if network_state["current_lora"] != repo_name:
-            network = create_network(
-                text_encoders,
-                pipe.unet,
-                lora_sd,
-                lora_weight,
-                device,
-            )
-            network_state["current_lora"] = repo_name
-            network_state["multiplier"] = lora_weight
-        elif network_state["multiplier"] != lora_weight:
-            network = create_network(
-                text_encoders,
-                pipe.unet,
-                lora_sd,
-                lora_weight,
-                device,
-            )
-            network_state["multiplier"] = lora_weight
-    else:
-        if network:
-            network.unapply_to()
-            network = None
-            network_state = {
-                "current_lora": None,
-                "multiplier": None,
-            }
+    width, height = utils.preprocess_image_dimensions(width, height)
 
     backup_scheduler = pipe.scheduler
-    pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)
+    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
 
     if use_upscaler:
         upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-
     metadata = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
@@ -344,11 +121,6 @@ def generate(
         "quality_tags": quality_selector,
     }
 
-    if use_lora:
-        metadata["use_lora"] = {"selected_lora": repo_name, "multiplier": lora_weight}
-    else:
-        metadata["use_lora"] = None
-
     if use_upscaler:
         new_width = int(width * upscale_by)
         new_height = int(height * upscale_by)
@@ -360,8 +132,7 @@ def generate(
         }
     else:
         metadata["use_upscaler"] = None
-
-    print(json.dumps(metadata, indent=4))
+    logger.info(json.dumps(metadata, indent=4))
 
     try:
         if use_upscaler:
@@ -375,8 +146,8 @@ def generate(
                 generator=generator,
                 output_type="latent",
             ).images
-            upscaled_latents = upscale(latents, "nearest-exact", upscale_by)
-            image = upscaler_pipe(
+            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+            images = upscaler_pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 image=upscaled_latents,
@@ -385,9 +156,9 @@ def generate(
                 strength=upscaler_strength,
                 generator=generator,
                 output_type="pil",
-            ).images[0]
+            ).images
         else:
-            image = pipe(
+            images = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 width=width,
@@ -396,194 +167,38 @@ def generate(
                 num_inference_steps=num_inference_steps,
                 generator=generator,
                 output_type="pil",
-            ).images[0]
-        if network:
-            network.unapply_to()
-            network = None
-        if profile is not None:
-            gr_user_history.save_image(
-                label=prompt,
-                image=image,
-                profile=profile,
-                metadata=metadata,
-            )
-        if image and IS_COLAB:
-            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
-            output_directory = "./outputs"
-            os.makedirs(output_directory, exist_ok=True)
-            filename = f"image_{current_time}.png"
-            filepath = os.path.join(output_directory, filename)
-
-            # Convert metadata to a string and save as a text chunk in the PNG
-            metadata_str = json.dumps(metadata)
-            info = PngImagePlugin.PngInfo()
-            info.add_text("metadata", metadata_str)
-            image.save(filepath, "PNG", pnginfo=info)
-            print(f"Image saved as {filepath} with metadata")
+            ).images
 
-        return image, metadata
+        if images and IS_COLAB:
+            for image in images:
+                filepath = utils.save_image(image, metadata, OUTPUT_DIR)
+                logger.info(f"Image saved as {filepath} with metadata")
 
+        return images, metadata
     except Exception as e:
-        print(f"An error occurred: {e}")
+        logger.exception(f"An error occurred: {e}")
         raise
     finally:
-        if network:
-            network.unapply_to()
-            network = None
-        if use_lora:
-            del lora_sd, text_encoders
         if use_upscaler:
             del upscaler_pipe
         pipe.scheduler = backup_scheduler
-        free_memory()
-
-
-examples = [
-    "1girl, arima kana, oshi no ko, solo, idol, idol clothes, one eye closed, red shirt, black skirt, black headwear, gloves, stage light, singing, open mouth, crowd, smile, pointing at viewer",
-    "1girl, c.c., code geass, white shirt, long sleeves, turtleneck, sitting, looking at viewer, eating, pizza, plate, fork, knife, table, chair, table, restaurant, cinematic angle, cinematic lighting",
-    "1girl, sakurauchi riko, \(love live\), queen hat, noble coat, red coat, noble shirt, sitting, crossed legs, gentle smile, parted lips, throne, cinematic angle",
-    "1girl, amiya \(arknights\), arknights, dirty face, outstretched hand, close-up, cinematic angle, foreshortening, dark, dark background",
-    "A boy and a girl, Emiya Shirou and Artoria Pendragon from fate series, having their breakfast in the dining room. Emiya Shirou wears white t-shirt and jacket. Artoria Pendragon wears white dress with blue neck ribbon. Rice, soup, and minced meats are served on the table. They look at each other while smiling happily",
-]
-
-quality_prompt_list = [
-    {
-        "name": "(None)",
-        "prompt": "{prompt}",
-        "negative_prompt": "nsfw, lowres, ",
-    },
-    {
-        "name": "Standard",
-        "prompt": "{prompt}, masterpiece, best quality",
-        "negative_prompt": "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, ",
-    },
-    {
-        "name": "Light",
-        "prompt": "{prompt}, (masterpiece), best quality, perfect face",
-        "negative_prompt": "nsfw, (low quality, worst quality:1.2), 3d, watermark, signature, ugly, poorly drawn, ",
-    },
-    {
-        "name": "Heavy",
-        "prompt": "{prompt}, (masterpiece), (best quality), (ultra-detailed), illustration, disheveled hair, perfect composition, moist skin, intricate details, earrings",
-        "negative_prompt": "nsfw, longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair, extra digit, fewer digits, cropped, worst quality, low quality, ",
-    },
-]
-
-sampler_list = [
-    "DPM++ 2M Karras",
-    "DPM++ SDE Karras",
-    "DPM++ 2M SDE Karras",
-    "Euler",
-    "Euler a",
-    "DDIM",
-]
-
-aspect_ratios = [
-    "1024 x 1024",
-    "1152 x 896",
-    "896 x 1152",
-    "1216 x 832",
-    "832 x 1216",
-    "1344 x 768",
-    "768 x 1344",
-    "1536 x 640",
-    "640 x 1536",
-    "Custom",
-]
-
-style_list = [
-    {
-        "name": "(None)",
-        "prompt": "{prompt}",
-        "negative_prompt": "",
-    },
-    {
-        "name": "Cinematic",
-        "prompt": "{prompt}, cinematic still, emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
-        "negative_prompt": "nsfw, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
-    },
-    {
-        "name": "Photographic",
-        "prompt": "{prompt}, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed",
-        "negative_prompt": "nsfw, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
-    },
-    {
-        "name": "Anime",
-        "prompt": "{prompt}, anime artwork, anime style, key visual, vibrant, studio anime,  highly detailed",
-        "negative_prompt": "nsfw, photo, deformed, black and white, realism, disfigured, low contrast",
-    },
-    {
-        "name": "Manga",
-        "prompt": "{prompt}, manga style, vibrant, high-energy, detailed, iconic, Japanese comic style",
-        "negative_prompt": "nsfw, ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
-    },
-    {
-        "name": "Digital Art",
-        "prompt": "{prompt}, concept art, digital artwork, illustrative, painterly, matte painting, highly detailed",
-        "negative_prompt": "nsfw, photo, photorealistic, realism, ugly",
-    },
-    {
-        "name": "Pixel art",
-        "prompt": "{prompt}, pixel-art, low-res, blocky, pixel art style, 8-bit graphics",
-        "negative_prompt": "nsfw, sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
-    },
-    {
-        "name": "Fantasy art",
-        "prompt": "{prompt}, ethereal fantasy concept art, magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
-        "negative_prompt": "nsfw, photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
-    },
-    {
-        "name": "Neonpunk",
-        "prompt": "{prompt}, neonpunk style, cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
-        "negative_prompt": "nsfw, painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
-    },
-    {
-        "name": "3D Model",
-        "prompt": "{prompt}, professional 3d model, octane render, highly detailed, volumetric, dramatic lighting",
-        "negative_prompt": "nsfw, ugly, deformed, noisy, low poly, blurry, painting",
-    },
-]
-
-thumbnail_cache = {}
+        utils.free_memory()
 
-with open("lora.toml", "r") as file:
-    data = toml.load(file)
 
-sdxl_loras = []
-saved_names = []
-for item in data["data"]:
-    model_path = hf_hub_download(item["repo"], item["weights"], token=HF_TOKEN)
-    saved_names.append(model_path)  # Store the path in saved_names
-
-    if model_path not in thumbnail_cache:
-        thumbnail_image = load_and_convert_thumbnail(model_path)
-        thumbnail_cache[model_path] = thumbnail_image
-    else:
-        thumbnail_image = thumbnail_cache[model_path]
-
-    sdxl_loras.append(
-        {
-            "image": thumbnail_image,  # Storing the PIL image object
-            "title": item["title"],
-            "repo": item["repo"],
-            "weights": item["weights"],
-            "multiplier": item.get("multiplier", "1.0"),
-        }
-    )
+if torch.cuda.is_available():
+    pipe = load_pipeline(MODEL)
+    logger.info("Loaded on Device!")
+else:
+    pipe = None
 
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list}
 quality_prompt = {
-    k["name"]: (k["prompt"], k["negative_prompt"]) for k in quality_prompt_list
+    k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.quality_prompt_list
 }
 
-# saved_names = [
-#     hf_hub_download(item["repo"], item["weights"], token=HF_TOKEN)
-#     for item in sdxl_loras
-# ]
-
-wildcard_files = load_wildcard_files("wildcard")
+wildcard_files = utils.load_wildcard_files("wildcard")
 
-with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
+with gr.Blocks(css="style.css") as demo:
     title = gr.HTML(
         f"""<h1><span>{DESCRIPTION}</span></h1>""",
         elem_id="title",
@@ -592,187 +207,131 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         f"""Gradio demo for [cagliostrolab/animagine-xl-3.0](https://huggingface.co/cagliostrolab/animagine-xl-3.0)""",
         elem_id="subtitle",
     )
-    gr.Markdown(
-        f"""Prompting is a bit different in this iteration, we train the model like this:
-        ```
-        1girl/1boy, character name, from what series, everything else in any order. 
-        ```
-        Prompting Tips
-        ```
-        1. Quality Tags: `masterpiece, best quality, high quality, normal quality, worst quality, low quality`
-        2. Year Tags: `oldest, early, mid, late, newest`
-        3. Rating tags: `rating: general, rating: sensitive, rating: questionable, rating: explicit, nsfw`
-        4. Escape character: `character name \(series\)`
-        5. Recommended settings: `Euler a, cfg 5-7, 25-28 steps`
-        6. It's recommended to use the exact danbooru tags for more accurate result
-        7. To use character wildcard, add this syntax to the prompt `__character__`.
-        ```
-        """,
-        elem_id="subtitle",
-    )   
     gr.DuplicateButton(
         value="Duplicate Space for private use",
         elem_id="duplicate-button",
         visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
     )
-    selected_state = gr.State()
-    with gr.Row():
-        with gr.Column(scale=2):
-            with gr.Tab("Txt2img"):
-                with gr.Group():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        max_lines=5,
-                        placeholder="Enter your prompt",
-                    )
-                    negative_prompt = gr.Text(
-                        label="Negative Prompt",
-                        max_lines=5,
-                        placeholder="Enter a negative prompt",
-                    )
-                    with gr.Accordion(label="Quality Tags", open=True):
-                        add_quality_tags = gr.Checkbox(label="Add Quality Tags", value=True)
-                        quality_selector = gr.Dropdown(
-                            label="Quality Tags Presets",
-                            interactive=True,
-                            choices=list(quality_prompt.keys()),
-                            value="Standard",
-                        )
-                    with gr.Row():
-                        use_lora = gr.Checkbox(label="Use LoRA", value=False)
-                with gr.Group(visible=False) as lora_group:
-                    selector_info = gr.Text(
-                        label="Selected LoRA",
-                        max_lines=1,
-                        value="No LoRA selected.",
-                    )
-                    lora_selection = gr.Gallery(
-                        value=[(item["image"], item["title"]) for item in sdxl_loras],
-                        label="Animagine XL 2.0 LoRA",
-                        show_label=False,
-                        columns=2,
-                        show_share_button=False,
-                    )
-                    lora_weight = gr.Slider(
-                        label="Multiplier",
-                        minimum=-2,
-                        maximum=2,
-                        step=0.05,
-                        value=1,
-                    )
-            with gr.Tab("Advanced Settings"):
-                with gr.Group():
-                    style_selector = gr.Radio(
-                        label="Style Preset",
-                        container=True,
-                        interactive=True,
-                        choices=list(styles.keys()),
-                        value="(None)",
-                    )
-                with gr.Group():
-                    aspect_ratio_selector = gr.Radio(
-                        label="Aspect Ratio",
-                        choices=aspect_ratios,
-                        value="896 x 1152",
-                        container=True,
-                    )
-                with gr.Group():
-                    use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
-                    with gr.Row() as upscaler_row:
-                        upscaler_strength = gr.Slider(
-                            label="Strength",
-                            minimum=0,
-                            maximum=1,
-                            step=0.05,
-                            value=0.55,
-                            visible=False,
-                        )
-                        upscale_by = gr.Slider(
-                            label="Upscale by",
-                            minimum=1,
-                            maximum=1.5,
-                            step=0.1,
-                            value=1.5,
-                            visible=False,
-                        )
-                with gr.Group(visible=False) as custom_resolution:
-                    with gr.Row():
-                        custom_width = gr.Slider(
-                            label="Width",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                        custom_height = gr.Slider(
-                            label="Height",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                with gr.Group():
-                    sampler = gr.Dropdown(
-                        label="Sampler",
-                        choices=sampler_list,
-                        interactive=True,
-                        value="Euler a",
-                    )
-                with gr.Group():
-                    seed = gr.Slider(
-                        label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0
-                    )
-
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                with gr.Group():
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=1,
-                            maximum=12,
-                            step=0.1,
-                            value=7.0,
-                        )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=28,
-                        )
-
-            with gr.Tab("Past Generation"):
-                gr_user_history.render()
-        with gr.Column(scale=3):
-            with gr.Blocks():
-                run_button = gr.Button("Generate", variant="primary")
-            result = gr.Image(label="Result", show_label=False)
-            with gr.Accordion(label="Generation Parameters", open=False):
-                gr_metadata = gr.JSON(label="Metadata", show_label=False)
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                outputs=[result, gr_metadata],
-                fn=generate,
-                cache_examples=CACHE_EXAMPLES,
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=5,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button(
+                "Generate", 
+                variant="primary", 
+                scale=0
+            )
+        result = gr.Gallery(
+            label="Result", 
+            columns=1, 
+            preview=True, 
+            show_label=False
+        )
+    with gr.Accordion(label="Advanced Settings", open=False):
+        negative_prompt = gr.Text(
+            label="Negative Prompt",
+            max_lines=5,
+            placeholder="Enter a negative prompt",
+        )
+        with gr.Row():
+            add_quality_tags = gr.Checkbox(
+                label="Add Quality Tags", 
+                value=True
+            )
+            quality_selector = gr.Dropdown(
+                label="Quality Tags Presets",
+                interactive=True,
+                choices=list(quality_prompt.keys()),
+                value="Standard",
+            )
+        style_selector = gr.Radio(
+            label="Style Preset",
+            container=True,
+            interactive=True,
+            choices=list(styles.keys()),
+            value="(None)",
+        )
+        aspect_ratio_selector = gr.Radio(
+            label="Aspect Ratio",
+            choices=config.aspect_ratios,
+            value="896 x 1152",
+            container=True,
+        )
+        with gr.Group(visible=False) as custom_resolution:
+            with gr.Row():
+                custom_width = gr.Slider(
+                    label="Width",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=8,
+                    value=1024,
+                )
+                custom_height = gr.Slider(
+                    label="Height",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=8,
+                    value=1024,
+                )
+        use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
+        with gr.Row() as upscaler_row:
+            upscaler_strength = gr.Slider(
+                label="Strength",
+                minimum=0,
+                maximum=1,
+                step=0.05,
+                value=0.55,
+                visible=False,
+            )
+            upscale_by = gr.Slider(
+                label="Upscale by",
+                minimum=1,
+                maximum=1.5,
+                step=0.1,
+                value=1.5,
+                visible=False,
             )
 
-    lora_selection.select(
-        update_selection,
-        outputs=[
-            selector_info,
-            selected_state,
-            lora_weight,
-        ],
-        queue=False,
-        show_progress=False,
-    )
-    use_lora.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_lora,
-        outputs=lora_group,
-        queue=False,
-        api_name=False,
+        sampler = gr.Dropdown(
+            label="Sampler",
+            choices=config.sampler_list,
+            interactive=True,
+            value="Euler a",
+        )
+        with gr.Row():
+            seed = gr.Slider(
+                label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
+            )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        with gr.Group():
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=1,
+                    maximum=12,
+                    step=0.1,
+                    value=7.0,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=28,
+                )
+    with gr.Accordion(label="Generation Parameters", open=False):
+        gr_metadata = gr.JSON(label="Metadata", show_label=False)
+    gr.Examples(
+        examples=config.examples,
+        inputs=prompt,
+        outputs=[result, gr_metadata],
+        fn=lambda *args, **kwargs: generate(*args, use_upscaler=True, **kwargs),
+        cache_examples=CACHE_EXAMPLES,
     )
     use_upscaler.change(
         fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
@@ -797,9 +356,6 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         custom_height,
         guidance_scale,
         num_inference_steps,
-        use_lora,
-        lora_weight,
-        selected_state,
         sampler,
         aspect_ratio_selector,
         style_selector,
@@ -807,11 +363,11 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         use_upscaler,
         upscaler_strength,
         upscale_by,
-        add_quality_tags
+        add_quality_tags,
     ]
 
     prompt.submit(
-        fn=randomize_seed_fn,
+        fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
         outputs=seed,
         queue=False,
@@ -823,7 +379,7 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         api_name="run",
     )
     negative_prompt.submit(
-        fn=randomize_seed_fn,
+        fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
         outputs=seed,
         queue=False,
@@ -835,7 +391,7 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         api_name=False,
     )
     run_button.click(
-        fn=randomize_seed_fn,
+        fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
         outputs=seed,
         queue=False,
@@ -846,4 +402,4 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         outputs=[result, gr_metadata],
         api_name=False,
     )
-demo.queue(max_size=30).launch(debug=IS_COLAB, share=IS_COLAB)
+demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)

config.py

@@ -0,0 +1,105 @@
+examples = [
+    "1girl, arima kana, oshi no ko, solo, idol, idol clothes, one eye closed, red shirt, black skirt, black headwear, gloves, stage light, singing, open mouth, crowd, smile, pointing at viewer",
+    "1girl, c.c., code geass, white shirt, long sleeves, turtleneck, sitting, looking at viewer, eating, pizza, plate, fork, knife, table, chair, table, restaurant, cinematic angle, cinematic lighting",
+    "1girl, sakurauchi riko, \(love live\), queen hat, noble coat, red coat, noble shirt, sitting, crossed legs, gentle smile, parted lips, throne, cinematic angle",
+    "1girl, amiya \(arknights\), arknights, dirty face, outstretched hand, close-up, cinematic angle, foreshortening, dark, dark background",
+    "A boy and a girl, Emiya Shirou and Artoria Pendragon from fate series, having their breakfast in the dining room. Emiya Shirou wears white t-shirt and jacket. Artoria Pendragon wears white dress with blue neck ribbon. Rice, soup, and minced meats are served on the table. They look at each other while smiling happily",
+]
+
+quality_prompt_list = [
+    {
+        "name": "(None)",
+        "prompt": "{prompt}",
+        "negative_prompt": "nsfw, lowres, ",
+    },
+    {
+        "name": "Standard",
+        "prompt": "{prompt}, masterpiece, best quality",
+        "negative_prompt": "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, ",
+    },
+    {
+        "name": "Light",
+        "prompt": "{prompt}, (masterpiece), best quality, perfect face",
+        "negative_prompt": "nsfw, (low quality, worst quality:1.2), 3d, watermark, signature, ugly, poorly drawn, ",
+    },
+    {
+        "name": "Heavy",
+        "prompt": "{prompt}, (masterpiece), (best quality), (ultra-detailed), illustration, disheveled hair, perfect composition, moist skin, intricate details, earrings",
+        "negative_prompt": "nsfw, longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair, extra digit, fewer digits, cropped, worst quality, low quality, ",
+    },
+]
+
+sampler_list = [
+    "DPM++ 2M Karras",
+    "DPM++ SDE Karras",
+    "DPM++ 2M SDE Karras",
+    "Euler",
+    "Euler a",
+    "DDIM",
+]
+
+aspect_ratios = [
+    "1024 x 1024",
+    "1152 x 896",
+    "896 x 1152",
+    "1216 x 832",
+    "832 x 1216",
+    "1344 x 768",
+    "768 x 1344",
+    "1536 x 640",
+    "640 x 1536",
+    "Custom",
+]
+
+style_list = [
+    {
+        "name": "(None)",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
+    {
+        "name": "Cinematic",
+        "prompt": "{prompt}, cinematic still, emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "nsfw, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+    {
+        "name": "Photographic",
+        "prompt": "{prompt}, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "nsfw, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
+    },
+    {
+        "name": "Anime",
+        "prompt": "{prompt}, anime artwork, anime style, key visual, vibrant, studio anime,  highly detailed",
+        "negative_prompt": "nsfw, photo, deformed, black and white, realism, disfigured, low contrast",
+    },
+    {
+        "name": "Manga",
+        "prompt": "{prompt}, manga style, vibrant, high-energy, detailed, iconic, Japanese comic style",
+        "negative_prompt": "nsfw, ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
+    },
+    {
+        "name": "Digital Art",
+        "prompt": "{prompt}, concept art, digital artwork, illustrative, painterly, matte painting, highly detailed",
+        "negative_prompt": "nsfw, photo, photorealistic, realism, ugly",
+    },
+    {
+        "name": "Pixel art",
+        "prompt": "{prompt}, pixel-art, low-res, blocky, pixel art style, 8-bit graphics",
+        "negative_prompt": "nsfw, sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
+    },
+    {
+        "name": "Fantasy art",
+        "prompt": "{prompt}, ethereal fantasy concept art, magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "negative_prompt": "nsfw, photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
+    },
+    {
+        "name": "Neonpunk",
+        "prompt": "{prompt}, neonpunk style, cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
+        "negative_prompt": "nsfw, painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "{prompt}, professional 3d model, octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "nsfw, ugly, deformed, noisy, low poly, blurry, painting",
+    },
+]

lora.toml

@@ -1,28 +0,0 @@
-[[data]]
-title = "Style Enhancer XL"
-repo = "Linaqruf/style-enhancer-xl-lora"
-weights = "style-enhancer-xl.safetensors"
-multiplier = 0.6
-[[data]]
-title = "Anime Detailer XL"
-repo = "Linaqruf/anime-detailer-xl-lora"
-weights = "anime-detailer-xl.safetensors"
-multiplier = 2.0
-
-[[data]]
-title = "Sketch Style XL"
-repo = "Linaqruf/sketch-style-xl-lora"
-weights = "sketch-style-xl.safetensors"
-multiplier = 0.6
-
-[[data]]
-title = "Pastel Style XL 2.0"
-repo = "Linaqruf/pastel-style-xl-lora"
-weights = "pastel-style-xl-v2.safetensors"
-multiplier = 0.6
-
-[[data]]
-title = "Anime Nouveau XL"
-repo = "Linaqruf/anime-nouveau-xl-lora"
-weights = "anime-nouveau-xl.safetensors"
-multiplier = 0.6

lora_diffusers.py

@@ -1,478 +0,0 @@
-"""
-LoRA module for Diffusers
-==========================
-
-This file works independently and is designed to operate with Diffusers.
-
-Credits
--------
-- Modified from: https://github.com/vladmandic/automatic/blob/master/modules/lora_diffusers.py
-- Originally from: https://github.com/kohya-ss/sd-scripts/blob/sdxl/networks/lora_diffusers.py
-"""
-
-import bisect
-import math
-import random
-from typing import Any, Dict, List, Mapping, Optional, Union
-from diffusers import UNet2DConditionModel
-import numpy as np
-from tqdm import tqdm
-from transformers import CLIPTextModel
-import torch
-
-
-def make_unet_conversion_map() -> Dict[str, str]:
-    unet_conversion_map_layer = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # loop over downblocks/upblocks
-        for j in range(2):
-            # loop over resnets/attentions for downblocks
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            if i < 3:
-                # no attention layers in down_blocks.3
-                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
-                sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
-
-        for j in range(3):
-            # loop over resnets/attentions for upblocks
-            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
-            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
-            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
-
-            # if i > 0: commentout for sdxl
-            # no attention layers in up_blocks.0
-            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
-            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
-            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
-
-        if i < 3:
-            # no downsample in down_blocks.3
-            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
-            sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
-
-            # no upsample in up_blocks.3
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
-            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    unet_conversion_map_resnet = [
-        # (stable-diffusion, HF Diffusers)
-        ("in_layers.0.", "norm1."),
-        ("in_layers.2.", "conv1."),
-        ("out_layers.0.", "norm2."),
-        ("out_layers.3.", "conv2."),
-        ("emb_layers.1.", "time_emb_proj."),
-        ("skip_connection.", "conv_shortcut."),
-    ]
-
-    unet_conversion_map = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    sd_hf_conversion_map = {sd.replace(".", "_")[:-1]: hf.replace(".", "_")[:-1] for sd, hf in unet_conversion_map}
-    return sd_hf_conversion_map
-
-
-UNET_CONVERSION_MAP = make_unet_conversion_map()
-
-
-class LoRAModule(torch.nn.Module):
-    """
-    replaces forward method of the original Linear, instead of replacing the original Linear module.
-    """
-
-    def __init__(
-        self,
-        lora_name,
-        org_module: torch.nn.Module,
-        multiplier=1.0,
-        lora_dim=4,
-        alpha=1,
-    ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
-
-        if org_module.__class__.__name__ == "Conv2d" or org_module.__class__.__name__ == "LoRACompatibleConv":
-            in_dim = org_module.in_channels
-            out_dim = org_module.out_channels
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-
-        self.lora_dim = lora_dim
-
-        if org_module.__class__.__name__ == "Conv2d" or org_module.__class__.__name__ == "LoRACompatibleConv":
-            kernel_size = org_module.kernel_size
-            stride = org_module.stride
-            padding = org_module.padding
-            self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
-            self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
-        else:
-            self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
-            self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
-
-        if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
-        alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
-        self.register_buffer("alpha", torch.tensor(alpha))  # 勾配計算に含めない / not included in gradient calculation
-
-        # same as microsoft's
-        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
-        torch.nn.init.zeros_(self.lora_up.weight)
-
-        self.multiplier = multiplier
-        self.org_module = [org_module]
-        self.enabled = True
-        self.network: LoRANetwork = None
-        self.org_forward = None
-
-    # override org_module's forward method
-    def apply_to(self, multiplier=None):
-        if multiplier is not None:
-            self.multiplier = multiplier
-        if self.org_forward is None:
-            self.org_forward = self.org_module[0].forward
-            self.org_module[0].forward = self.forward
-
-    # restore org_module's forward method
-    def unapply_to(self):
-        if self.org_forward is not None:
-            self.org_module[0].forward = self.org_forward
-
-    # forward with lora
-    # scale is used LoRACompatibleConv, but we ignore it because we have multiplier
-    def forward(self, x, scale=1.0):
-        if not self.enabled:
-            return self.org_forward(x)
-        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
-
-    def set_network(self, network):
-        self.network = network
-
-    # merge lora weight to org weight
-    def merge_to(self, multiplier=1.0):
-        # get lora weight
-        lora_weight = self.get_weight(multiplier)
-
-        # get org weight
-        org_sd = self.org_module[0].state_dict()
-        org_weight = org_sd["weight"]
-        weight = org_weight + lora_weight.to(org_weight.device, dtype=org_weight.dtype)
-
-        # set weight to org_module
-        org_sd["weight"] = weight
-        self.org_module[0].load_state_dict(org_sd)
-
-    # restore org weight from lora weight
-    def restore_from(self, multiplier=1.0):
-        # get lora weight
-        lora_weight = self.get_weight(multiplier)
-
-        # get org weight
-        org_sd = self.org_module[0].state_dict()
-        org_weight = org_sd["weight"]
-        weight = org_weight - lora_weight.to(org_weight.device, dtype=org_weight.dtype)
-
-        # set weight to org_module
-        org_sd["weight"] = weight
-        self.org_module[0].load_state_dict(org_sd)
-
-    # return lora weight
-    def get_weight(self, multiplier=None):
-        if multiplier is None:
-            multiplier = self.multiplier
-
-        # get up/down weight from module
-        up_weight = self.lora_up.weight.to(torch.float)
-        down_weight = self.lora_down.weight.to(torch.float)
-
-        # pre-calculated weight
-        if len(down_weight.size()) == 2:
-            # linear
-            weight = self.multiplier * (up_weight @ down_weight) * self.scale
-        elif down_weight.size()[2:4] == (1, 1):
-            # conv2d 1x1
-            weight = (
-                self.multiplier
-                * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
-                * self.scale
-            )
-        else:
-            # conv2d 3x3
-            conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-            weight = self.multiplier * conved * self.scale
-
-        return weight
-
-
-# Create network from weights for inference, weights are not loaded here
-def create_network_from_weights(
-    text_encoder: Union[CLIPTextModel, List[CLIPTextModel]], unet: UNet2DConditionModel, weights_sd: Dict, multiplier: float = 1.0
-):
-    # get dim/alpha mapping
-    modules_dim = {}
-    modules_alpha = {}
-    for key, value in weights_sd.items():
-        if "." not in key:
-            continue
-
-        lora_name = key.split(".")[0]
-        if "alpha" in key:
-            modules_alpha[lora_name] = value
-        elif "lora_down" in key:
-            dim = value.size()[0]
-            modules_dim[lora_name] = dim
-            # print(lora_name, value.size(), dim)
-
-    # support old LoRA without alpha
-    for key in modules_dim.keys():
-        if key not in modules_alpha:
-            modules_alpha[key] = modules_dim[key]
-
-    return LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha)
-
-
-def merge_lora_weights(pipe, weights_sd: Dict, multiplier: float = 1.0):
-    text_encoders = [pipe.text_encoder, pipe.text_encoder_2] if hasattr(pipe, "text_encoder_2") else [pipe.text_encoder]
-    unet = pipe.unet
-
-    lora_network = create_network_from_weights(text_encoders, unet, weights_sd, multiplier=multiplier)
-    lora_network.load_state_dict(weights_sd)
-    lora_network.merge_to(multiplier=multiplier)
-
-
-# block weightや学習に対応しない簡易版 / simple version without block weight and training
-class LoRANetwork(torch.nn.Module):
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
-    UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = "lora_unet"
-    LORA_PREFIX_TEXT_ENCODER = "lora_te"
-
-    # SDXL: must starts with LORA_PREFIX_TEXT_ENCODER
-    LORA_PREFIX_TEXT_ENCODER1 = "lora_te1"
-    LORA_PREFIX_TEXT_ENCODER2 = "lora_te2"
-
-    def __init__(
-        self,
-        text_encoder: Union[List[CLIPTextModel], CLIPTextModel],
-        unet: UNet2DConditionModel,
-        multiplier: float = 1.0,
-        modules_dim: Optional[Dict[str, int]] = None,
-        modules_alpha: Optional[Dict[str, int]] = None,
-        varbose: Optional[bool] = False,
-    ) -> None:
-        super().__init__()
-        self.multiplier = multiplier
-
-        print(f"create LoRA network from weights")
-
-        # convert SDXL Stability AI's U-Net modules to Diffusers
-        converted = self.convert_unet_modules(modules_dim, modules_alpha)
-        if converted:
-            print(f"converted {converted} Stability AI's U-Net LoRA modules to Diffusers (SDXL)")
-
-        # create module instances
-        def create_modules(
-            is_unet: bool,
-            text_encoder_idx: Optional[int],  # None, 1, 2
-            root_module: torch.nn.Module,
-            target_replace_modules: List[torch.nn.Module],
-        ) -> List[LoRAModule]:
-            prefix = (
-                self.LORA_PREFIX_UNET
-                if is_unet
-                else (
-                    self.LORA_PREFIX_TEXT_ENCODER
-                    if text_encoder_idx is None
-                    else (self.LORA_PREFIX_TEXT_ENCODER1 if text_encoder_idx == 1 else self.LORA_PREFIX_TEXT_ENCODER2)
-                )
-            )
-            loras = []
-            skipped = []
-            for name, module in root_module.named_modules():
-                if module.__class__.__name__ in target_replace_modules:
-                    for child_name, child_module in module.named_modules():
-                        is_linear = (
-                            child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "LoRACompatibleLinear"
-                        )
-                        is_conv2d = (
-                            child_module.__class__.__name__ == "Conv2d" or child_module.__class__.__name__ == "LoRACompatibleConv"
-                        )
-
-                        if is_linear or is_conv2d:
-                            lora_name = prefix + "." + name + "." + child_name
-                            lora_name = lora_name.replace(".", "_")
-
-                            if lora_name not in modules_dim:
-                                # print(f"skipped {lora_name} (not found in modules_dim)")
-                                skipped.append(lora_name)
-                                continue
-
-                            dim = modules_dim[lora_name]
-                            alpha = modules_alpha[lora_name]
-                            lora = LoRAModule(
-                                lora_name,
-                                child_module,
-                                self.multiplier,
-                                dim,
-                                alpha,
-                            )
-                            loras.append(lora)
-            return loras, skipped
-
-        text_encoders = text_encoder if type(text_encoder) == list else [text_encoder]
-
-        # create LoRA for text encoder
-        # 毎回すべてのモジュールを作るのは無駄なので要検討 / it is wasteful to create all modules every time, need to consider
-        self.text_encoder_loras: List[LoRAModule] = []
-        skipped_te = []
-        for i, text_encoder in enumerate(text_encoders):
-            if len(text_encoders) > 1:
-                index = i + 1
-            else:
-                index = None
-
-            text_encoder_loras, skipped = create_modules(False, index, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-            self.text_encoder_loras.extend(text_encoder_loras)
-            skipped_te += skipped
-        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
-        if len(skipped_te) > 0:
-            print(f"skipped {len(skipped_te)} modules because of missing weight for text encoder.")
-
-        # extend U-Net target modules to include Conv2d 3x3
-        target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE + LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
-
-        self.unet_loras: List[LoRAModule]
-        self.unet_loras, skipped_un = create_modules(True, None, unet, target_modules)
-        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
-        if len(skipped_un) > 0:
-            print(f"skipped {len(skipped_un)} modules because of missing weight for U-Net.")
-
-        # assertion
-        names = set()
-        for lora in self.text_encoder_loras + self.unet_loras:
-            names.add(lora.lora_name)
-        for lora_name in modules_dim.keys():
-            assert lora_name in names, f"{lora_name} is not found in created LoRA modules."
-
-        # make to work load_state_dict
-        for lora in self.text_encoder_loras + self.unet_loras:
-            self.add_module(lora.lora_name, lora)
-
-    # SDXL: convert SDXL Stability AI's U-Net modules to Diffusers
-    def convert_unet_modules(self, modules_dim, modules_alpha):
-        converted_count = 0
-        not_converted_count = 0
-
-        map_keys = list(UNET_CONVERSION_MAP.keys())
-        map_keys.sort()
-
-        for key in list(modules_dim.keys()):
-            if key.startswith(LoRANetwork.LORA_PREFIX_UNET + "_"):
-                search_key = key.replace(LoRANetwork.LORA_PREFIX_UNET + "_", "")
-                position = bisect.bisect_right(map_keys, search_key)
-                map_key = map_keys[position - 1]
-                if search_key.startswith(map_key):
-                    new_key = key.replace(map_key, UNET_CONVERSION_MAP[map_key])
-                    modules_dim[new_key] = modules_dim[key]
-                    modules_alpha[new_key] = modules_alpha[key]
-                    del modules_dim[key]
-                    del modules_alpha[key]
-                    converted_count += 1
-                else:
-                    not_converted_count += 1
-        assert (
-            converted_count == 0 or not_converted_count == 0
-        ), f"some modules are not converted: {converted_count} converted, {not_converted_count} not converted"
-        return converted_count
-
-    def set_multiplier(self, multiplier):
-        self.multiplier = multiplier
-        for lora in self.text_encoder_loras + self.unet_loras:
-            lora.multiplier = self.multiplier
-
-    def apply_to(self, multiplier=1.0, apply_text_encoder=True, apply_unet=True):
-        if apply_text_encoder:
-            print("enable LoRA for text encoder")
-            for lora in self.text_encoder_loras:
-                lora.apply_to(multiplier)
-        if apply_unet:
-            print("enable LoRA for U-Net")
-            for lora in self.unet_loras:
-                lora.apply_to(multiplier)
-
-    def unapply_to(self):
-        for lora in self.text_encoder_loras + self.unet_loras:
-            lora.unapply_to()
-
-    def merge_to(self, multiplier=1.0):
-        print("merge LoRA weights to original weights")
-        for lora in tqdm(self.text_encoder_loras + self.unet_loras):
-            lora.merge_to(multiplier)
-        print(f"weights are merged")
-
-    def restore_from(self, multiplier=1.0):
-        print("restore LoRA weights from original weights")
-        for lora in tqdm(self.text_encoder_loras + self.unet_loras):
-            lora.restore_from(multiplier)
-        print(f"weights are restored")
-
-    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
-        # convert SDXL Stability AI's state dict to Diffusers' based state dict
-        map_keys = list(UNET_CONVERSION_MAP.keys())  # prefix of U-Net modules
-        map_keys.sort()
-        for key in list(state_dict.keys()):
-            if key.startswith(LoRANetwork.LORA_PREFIX_UNET + "_"):
-                search_key = key.replace(LoRANetwork.LORA_PREFIX_UNET + "_", "")
-                position = bisect.bisect_right(map_keys, search_key)
-                map_key = map_keys[position - 1]
-                if search_key.startswith(map_key):
-                    new_key = key.replace(map_key, UNET_CONVERSION_MAP[map_key])
-                    state_dict[new_key] = state_dict[key]
-                    del state_dict[key]
-
-        # in case of V2, some weights have different shape, so we need to convert them
-        # because V2 LoRA is based on U-Net created by use_linear_projection=False
-        my_state_dict = self.state_dict()
-        for key in state_dict.keys():
-            if state_dict[key].size() != my_state_dict[key].size():
-                # print(f"convert {key} from {state_dict[key].size()} to {my_state_dict[key].size()}")
-                state_dict[key] = state_dict[key].view(my_state_dict[key].size())
-
-        return super().load_state_dict(state_dict, strict)

requirements.txt

@@ -1,11 +1,10 @@
-accelerate==0.24.1
-diffusers==0.23.0
-gradio==4.2.0
+accelerate==0.27.2
+diffusers==0.26.3
+gradio==4.20.0
 invisible-watermark==0.2.0
-Pillow==10.1.0
+Pillow==10.2.0
+spaces==0.24.0
 torch==2.0.1
-transformers==4.35.0
-toml==0.10.2
+transformers==4.38.1
 omegaconf==2.3.0
 timm==0.9.10
-git+https://huggingface.co/spaces/Wauplin/gradio-user-history

style.css

@@ -1,11 +1,6 @@
 h1 {
   text-align: center;
-  font-size: 10vw; /* relative to the viewport width */
-}
-
-h2 {
-  text-align: center;
-  font-size: 10vw; /* relative to the viewport width */
+  display: block;
 }
 
 #duplicate-button {
@@ -15,24 +10,12 @@ h2 {
   border-radius: 100vh;
 }
 
-#component-0 {
-  max-width: 80%; /* relative to the parent element's width */
+.gradio-container {
+  max-width: 730px !important;
   margin: auto;
   padding-top: 1.5rem;
 }
 
-/* You can also use media queries to adjust your style for different screen sizes */
-@media (max-width: 600px) {
-  #component-0 {
-    max-width: 90%;
-    padding-top: 1rem;
-  }
-}
-
-#gallery .grid-wrap{
-  min-height: 25%;
-}
-
 #title-container {
     display: flex;
     justify-content: center;
@@ -43,18 +26,9 @@ h2 {
 #title {
     font-size: 3em;
     text-align: center;
-    color: #333;
-    font-family: 'Helvetica Neue', sans-serif;
-    text-transform: uppercase;
     background: transparent;
   }
 
-#title span {
-    background: -webkit-linear-gradient(45deg, #4EACEF, #28b485);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-}
-
 #subtitle {
   text-align: center;
-}
+}

utils.py

@@ -1,7 +1,179 @@
+import gc
+import os
+import random
+import numpy as np
+import json
+import torch
+from PIL import Image, PngImagePlugin
+from datetime import datetime
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+from diffusers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+@dataclass
+class StyleConfig:
+    prompt: str
+    negative_prompt: str
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+def seed_everything(seed: int) -> torch.Generator:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return generator
+
+
+def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]:
+    if aspect_ratio == "Custom":
+        return None
+    width, height = aspect_ratio.split(" x ")
+    return int(width), int(height)
+
+
+def aspect_ratio_handler(
+    aspect_ratio: str, custom_width: int, custom_height: int
+) -> Tuple[int, int]:
+    if aspect_ratio == "Custom":
+        return custom_width, custom_height
+    else:
+        width, height = parse_aspect_ratio(aspect_ratio)
+        return width, height
+
+
+def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
+    scheduler_factory_map = {
+        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
+        ),
+        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
+        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
+            scheduler_config
+        ),
+        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
+    }
+    return scheduler_factory_map.get(name, lambda: None)()
+
+
+def free_memory() -> None:
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def preprocess_prompt(
+    style_dict,
+    style_name: str,
+    positive: str,
+    negative: str = "",
+    add_style: bool = True,
+) -> Tuple[str, str]:
+    p, n = style_dict.get(style_name, style_dict["(None)"])
+
+    if add_style and positive.strip():
+        formatted_positive = p.format(prompt=positive)
+    else:
+        formatted_positive = positive
+
+    combined_negative = n
+    if negative.strip():
+        if combined_negative:
+            combined_negative += ", " + negative
+        else:
+            combined_negative = negative
+
+    return formatted_positive, combined_negative
+
+
+def common_upscale(
+    samples: torch.Tensor,
+    width: int,
+    height: int,
+    upscale_method: str,
+) -> torch.Tensor:
+    return torch.nn.functional.interpolate(
+        samples, size=(height, width), mode=upscale_method
+    )
+
+
+def upscale(
+    samples: torch.Tensor, upscale_method: str, scale_by: float
+) -> torch.Tensor:
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+    return common_upscale(samples, width, height, upscale_method)
+
+
+def load_wildcard_files(wildcard_dir: str) -> Dict[str, str]:
+    wildcard_files = {}
+    for file in os.listdir(wildcard_dir):
+        if file.endswith(".txt"):
+            key = f"__{file.split('.')[0]}__"  # Create a key like __character__
+            wildcard_files[key] = os.path.join(wildcard_dir, file)
+    return wildcard_files
+
+
+def get_random_line_from_file(file_path: str) -> str:
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+        if not lines:
+            return ""
+        return random.choice(lines).strip()
+
+
+def add_wildcard(prompt: str, wildcard_files: Dict[str, str]) -> str:
+    for key, file_path in wildcard_files.items():
+        if key in prompt:
+            wildcard_line = get_random_line_from_file(file_path)
+            prompt = prompt.replace(key, wildcard_line)
+    return prompt
+
+
+def preprocess_image_dimensions(width, height):
+    if width % 8 != 0:
+        width = width - (width % 8)
+    if height % 8 != 0:
+        height = height - (height % 8)
+    return width, height
+
+
+def save_image(image, metadata, output_dir):
+    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+    os.makedirs(output_dir, exist_ok=True)
+    filename = f"image_{current_time}.png"
+    filepath = os.path.join(output_dir, filename)
+
+    metadata_str = json.dumps(metadata)
+    info = PngImagePlugin.PngInfo()
+    info.add_text("metadata", metadata_str)
+    image.save(filepath, "PNG", pnginfo=info)
+    return filepath
+    
+    
 def is_google_colab():
     try:
         import google.colab
-
         return True
     except:
         return False