app.py

import gradio as gr
from pathlib import Path


loaded_models = {}
model_info_dict = {}


def list_sub(a, b):
    return [e for e in a if e not in b]


def list_uniq(l):
        return sorted(set(l), key=l.index)


def is_repo_name(s):
    import re
    return re.fullmatch(r'^[^/]+?/[^/]+?$', s)


def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30):
    from huggingface_hub import HfApi
    api = HfApi()
    default_tags = ["diffusers"]
    models = []
    try:
        model_infos = api.list_models(author=author, task="text-to-image", pipeline_tag="text-to-image",
                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
    except Exception as e:
        print(f"Error: Failed to list models.")
        print(e)
        return models
    for model in model_infos:
        if not model.private and not model.gated:
           if not_tag and not_tag in model.tags: continue
           models.append(model.id)
           if len(models) == limit: break
    return models


models = find_model_list("John6666", ["anime"])


def get_t2i_model_info_dict(repo_id: str):
    from huggingface_hub import HfApi
    api = HfApi()
    info = {"md": "None"}
    try:
        if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info
        model = api.model_info(repo_id=repo_id)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info.")
        print(e)
        return info
    if model.private or model.gated: return info
    try:
        tags = model.tags
    except Exception:
        return info
    if not 'diffusers' in model.tags: return info
    if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
    elif 'diffusers:StableDiffusionPipeline' in tags: info["ver"] = "SD1.5"
    elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
    else: info["ver"] = "Other"
    info["url"] = f"https://huggingface.co/{repo_id}/"
    if model.card_data and model.card_data.tags:
        info["tags"] = model.card_data.tags
    info["downloads"] = model.downloads
    info["likes"] = model.likes
    info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
    un_tags = ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']
    descs = [info["ver"]] + list_sub(info["tags"], un_tags) + [f'DLs: {info["downloads"]}'] + [f'❤: {info["likes"]}'] + [info["last_modified"]]
    info["md"] = f'Model Info: {", ".join(descs)} [Model Repo]({info["url"]})'
    return info


def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
    from datetime import datetime, timezone, timedelta
    progress(0, desc="Updating gallery...")
    dt_now = datetime.now(timezone(timedelta(hours=9)))
    basename = dt_now.strftime('%Y%m%d_%H%M%S_')
    i = 1
    if not images: return images
    output_images = []
    output_paths = []
    for image in images:
        filename = f'{image[1]}_{basename}{str(i)}.png'
        i += 1
        oldpath = Path(image[0])
        newpath = oldpath
        try:
            if oldpath.stem == "image" and oldpath.exists():
                newpath = oldpath.resolve().rename(Path(filename).resolve())
        except Exception as e:
           print(e)
           pass
        finally:
            output_paths.append(str(newpath))
            output_images.append((str(newpath), str(filename)))
    progress(1, desc="Gallery updated.")
    return gr.update(value=output_images), gr.update(value=output_paths)


def load_model(model_name: str):
    if model_name in loaded_models.keys(): return loaded_models[model_name]
    try:
        loaded_models[model_name] = gr.load(f'models/{model_name}')
        print(f"Loaded: {model_name}")
    except Exception as e:
        if model_name in loaded_models.keys(): del loaded_models[model_name]
        print(f"Failed to load: {model_name}")
        print(e)
        return None
    try:
        model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
    except Exception as e:
        if model_name in model_info_dict.keys(): del model_info_dict[model_name]
        print(e)
    return loaded_models[model_name]


for model in models:
    load_model(model)


def get_model_info_md(model_name: str):
    if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "")


def change_model(model_name: str):
    load_model(model_name)
    return get_model_info_md(model_name)


def infer(prompt: str, model_name: str, recom_prompt: bool, progress=gr.Progress(track_tqdm=True)):
    from PIL import Image
    import random
    seed = ""
    rand = random.randint(1, 500)
    for i in range(rand):
        seed += " "
    rprompt = ", highly detailed, masterpiece, best quality, very aesthetic, absurdres, " if recom_prompt else ""
    caption = model_name.split("/")[-1]
    try:
        model = load_model(model_name)
        if not model: return (Image(), None)
        image_path = model(prompt + rprompt + seed)
        image = Image.open(image_path).convert('RGB')
    except Exception as e:
        print(e)
        return (Image(), None)
    return (image, caption)


def infer_multi(prompt: str, model_name: str, recom_prompt: bool, image_num: float, results: list, progress=gr.Progress(track_tqdm=True)):
    image_num = int(image_num)
    images = results if results else []
    for i in range(image_num):
        images.append(infer(prompt, model_name, recom_prompt))
        yield images


css = """"""

with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
    with gr.Column(): 
        model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
        model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]))
        image_num = gr.Slider(label="Number of Images", minimum=1, maximum=8, value=1, step=1)
        recom_prompt = gr.Checkbox(label="Recommended Prompt", value=True)
        prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...")
        run_button = gr.Button("Generate Image")
        results = gr.Gallery(label="Gallery", interactive=False, show_download_button=True, show_share_button=False,
                              container=True, format="png", object_fit="contain")
        image_files = gr.Files(label="Download", interactive=False)
        clear_results = gr.Button("Clear Gallery and Download")
    gr.Markdown(
        f"""This demo was created in reference to the following demos.
- [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood).
- [Yntec/ToyWorldXL](https://huggingface.co/spaces/Yntec/ToyWorldXL).
            """
    )
    gr.DuplicateButton(value="Duplicate Space")

    model_name.change(change_model, [model_name], [model_info], queue=False, show_api=False)
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer_multi,
        inputs=[prompt, model_name, recom_prompt, image_num, results],
        outputs=[results],
        queue=True,
        show_progress="full",
        show_api=True,
    ).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
    clear_results.click(lambda: (None, None), None, [results, image_files], queue=False, show_api=False)

demo.queue()
demo.launch()