File size: 3,042 Bytes
7946a9d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import os
import json
from typing import Union, List
from PIL import Image
def get_examples(examples_dir: Union[str, List[str]] = None, use_lora: bool = None) -> List:
directories = [examples_dir] if isinstance(examples_dir, str) else examples_dir or []
valid_dirs = [d for d in directories if os.path.isdir(d)]
if not valid_dirs:
return get_provided_examples(use_lora)
examples = []
for dir_path in valid_dirs:
for subdir in sorted(os.path.join(dir_path, d) for d in os.listdir(dir_path) if os.path.isdir(os.path.join(dir_path, d))):
config_path = os.path.join(subdir, "config.json")
image_path = os.path.join(subdir, "result.png")
if not (os.path.isfile(config_path) and os.path.isfile(image_path)):
continue
with open(config_path, 'r') as f:
config = json.load(f)
required_keys = ["prompt", "height", "width", "num_inference_steps", "guidance_scale", "seed", "image"]
if config.get("use_lora", False):
required_keys.extend(["lora_model_id", "base_model_id",
# "lora_rank",
"lora_scale"])
else:
required_keys.append("finetune_model_id")
if set(required_keys) - set(config.keys()) or config["image"] != "result.png":
continue
try:
image = Image.open(image_path)
except Exception:
continue
if use_lora is not None and config.get("use_lora", False) != use_lora:
continue
example = [config["prompt"], config["height"], config["width"], config["num_inference_steps"],
config["guidance_scale"], config["seed"], image]
example.extend([config["lora_model_id"], config["base_model_id"],
# config["lora_rank"],
config["lora_scale"]]
if config.get("use_lora", False) else [config["finetune_model_id"]])
examples.append(example)
return examples or get_provided_examples(use_lora)
def get_provided_examples(use_lora: bool = False) -> list:
example_path = f"apps/gradio_app/assets/examples/Ghibli-Stable-Diffusion-2.1-{'LoRA' if use_lora else 'Base-finetuning'}/1/result.png"
image = Image.open(example_path) if os.path.exists(example_path) else None
return [[
"a cat is laying on a sofa in Ghibli style" if use_lora else "a serene landscape in Ghibli style",
512, 768 if use_lora else 512, 100 if use_lora else 50, 10.0 if use_lora else 3.5, 789 if use_lora else 42,
image, "danhtran2mind/Ghibli-Stable-Diffusion-2.1-LoRA" if use_lora else "danhtran2mind/Ghibli-Stable-Diffusion-2.1-Base-finetuning",
"stabilityai/stable-diffusion-2-1" if use_lora else None, 64 if use_lora else None, 0.9 if use_lora else None
]] |