t2i_space.py

from pathlib import Path
import os


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


def is_repo_exists(repo_id):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        if api.repo_exists(repo_id=repo_id, repo_type="model"): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}. ")
        return True # for safe


def is_repo_t2i(repo_id):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        model_info =  api.repo_info(repo_id=repo_id, repo_type="model")
        if model_info.pipeline_tag == "text-to-image": return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}. ")
        return True # for safe
    

def save_space_contents(repo_id: str, dir: str):
    if not is_repo_name(repo_id) or not is_repo_exists(repo_id) or not is_repo_t2i(repo_id):
        print(f"Error: Invalid repo ID: {repo_id}. ")
        return []
    os.makedirs(dir, exist_ok=True)
    model_name = repo_id.split("/")[-1]
    app_py = f"""import gradio as gr
demo = gr.load("{repo_id}", src="models").launch()
demo.launch()
"""
    readme_md = f"""---
title: {model_name} Text-to-Image
emoji: 🖼
colorFrom: purple
colorTo: red
sdk: gradio
sdk_version: 4.38.1
app_file: app.py
pinned: false
short_description: {repo_id} | Text-to-Image
---

Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

"""
    app_path = str(Path(dir, "app.py"))
    with open(app_path, mode='w', encoding="utf-8") as f:
        f.write(app_py)
    readme_path = str(Path(dir, "README.md"))
    with open(readme_path, mode='w', encoding="utf-8") as f:
        f.write(readme_md)
    return [app_path, readme_path]


def get_t2i_space_contents(repo_id: str):
    return save_space_contents(repo_id, "./temp/")