|
|
import spaces
|
|
|
import gradio as gr
|
|
|
from huggingface_hub import snapshot_download
|
|
|
import os, platform, shutil
|
|
|
from pathlib import Path
|
|
|
from importlib import metadata as md
|
|
|
from typing import List
|
|
|
|
|
|
print("Python:", platform.python_version())
|
|
|
print("OS:", platform.uname().system)
|
|
|
print("\n".join(sorted(f"{d.metadata['Name']}=={d.version}" for d in md.distributions())))
|
|
|
|
|
|
repo_id = "stabilityai/stable-diffusion-3.5-medium"
|
|
|
local_dir="./test"
|
|
|
cache_dir="./test-cache"
|
|
|
ALLOW_PATTERNS = ["transformer/diffusion_pytorch_model*", "vae/diffusion_pytorch_model*", "text_encoder/model*", "text_encoder_2/model*", "text_encoder_3/model*", "text_encoder/pytorch_model*", "text_encoder_2/pytorch_model*", "text_encoder_3/pytorch_model*"]
|
|
|
IGNORE_PATTERNS = ["*fp16*", "**/*fp16*"]
|
|
|
IGNORE_PATTERNS_EX = ["**/model*", "**/*.safetensors"]
|
|
|
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN", None)
|
|
|
|
|
|
@spaces.GPU
|
|
|
def test(repo_id: str, allow: List[str], ignore: List[str], del_cache: bool, progress=gr.Progress(track_tqdm=True)) -> str:
|
|
|
if del_cache: shutil.rmtree(cache_dir, ignore_errors=True)
|
|
|
shutil.rmtree(local_dir, ignore_errors=True)
|
|
|
kwargs = {}
|
|
|
if len(allow) > 0: kwargs["allow_patterns"] = allow
|
|
|
if len(ignore) > 0: kwargs["ignore_patterns"] = ignore
|
|
|
snapshot_download(
|
|
|
repo_id=repo_id,
|
|
|
local_dir=local_dir,
|
|
|
cache_dir=cache_dir,
|
|
|
token=HF_TOKEN,
|
|
|
**kwargs,
|
|
|
)
|
|
|
result = (f"Python:{platform.python_version()}\n"
|
|
|
f"OS:{platform.uname().system}\n\n"
|
|
|
f"repo_id:{repo_id}\n"
|
|
|
f"allow_patterns:{", ".join(allow)}\n"
|
|
|
f"ignore_patterns:{", ".join(ignore)}\n\n"
|
|
|
f'Downloaded:\n{"\n".join([str(p) for p in Path(local_dir).rglob("*") if p.is_file()])}')
|
|
|
return result
|
|
|
|
|
|
with gr.Blocks() as demo:
|
|
|
repo_id = gr.Textbox(label="Repo ID", value=repo_id, lines=1, visible=False)
|
|
|
allow = gr.CheckboxGroup(label="allow_patterns", choices=ALLOW_PATTERNS, value=ALLOW_PATTERNS)
|
|
|
ignore = gr.CheckboxGroup(label="ignore_patterns", choices=IGNORE_PATTERNS + IGNORE_PATTERNS_EX, value=IGNORE_PATTERNS)
|
|
|
del_cache = gr.Checkbox(label="Delete cache", value=False)
|
|
|
run_button = gr.Button("Download", variant="primary")
|
|
|
info = gr.Textbox(label="Result", value="", show_copy_button=True)
|
|
|
|
|
|
run_button.click(test, [repo_id, allow, ignore, del_cache], [info])
|
|
|
|
|
|
demo.queue().launch()
|
|
|
|
