Upload folder using huggingface_hub
Browse files- app.py +157 -0
- requirements.txt +5 -0
app.py
ADDED
@@ -0,0 +1,157 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os, re, json, shutil, tempfile
|
2 |
+
from typing import Dict, Any, List, Optional, Tuple
|
3 |
+
import requests
|
4 |
+
import gradio as gr
|
5 |
+
|
6 |
+
try:
|
7 |
+
from huggingface_hub import HfApi, create_repo, upload_folder, whoami
|
8 |
+
except Exception as e:
|
9 |
+
HfApi = None # Will error at runtime with a clear message
|
10 |
+
|
11 |
+
CIVITAI_API_TOKEN = os.getenv("CIVITAI_API_TOKEN", "").strip()
|
12 |
+
COOKIE_INFO = os.getenv("COOKIE_INFO", "").strip()
|
13 |
+
|
14 |
+
BASE_MODEL_MAP: Dict[str, str] = {
|
15 |
+
"sd 1.4": "CompVis/stable-diffusion-v1-4",
|
16 |
+
"sd 1.5": "runwayml/stable-diffusion-v1-5",
|
17 |
+
"sd 2.1": "stabilityai/stable-diffusion-2-1",
|
18 |
+
"sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
|
19 |
+
"sdxl 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
|
20 |
+
"flux.1 dev": "black-forest-labs/FLUX.1-dev",
|
21 |
+
"flux.1 schnell": "black-forest-labs/FLUX.1-schnell",
|
22 |
+
"flux": "black-forest-labs/FLUX.1-dev",
|
23 |
+
}
|
24 |
+
|
25 |
+
LICENSE_ALLOWLIST = {
|
26 |
+
"cc0-1.0", "cc0", "mit", "apache-2.0", "bsd-3-clause", "bsd-2-clause",
|
27 |
+
"openrail", "openrail+", "openrail++", "bigscience-openrail-m",
|
28 |
+
}
|
29 |
+
|
30 |
+
def _slugify(s: str, maxlen: int = 64) -> str:
|
31 |
+
s = s.strip().lower()
|
32 |
+
s = re.sub(r"[^a-z0-9\-_]+", "-", s)
|
33 |
+
s = re.sub(r"-{2,}", "-", s).strip("-")
|
34 |
+
if not s:
|
35 |
+
s = "model"
|
36 |
+
return s[:maxlen]
|
37 |
+
|
38 |
+
def http_get(url: str, stream: bool = False):
|
39 |
+
headers = {"User-Agent": "HF-Space-CivitAI-Importer/1.0"}
|
40 |
+
if COOKIE_INFO:
|
41 |
+
headers["Cookie"] = COOKIE_INFO
|
42 |
+
if CIVITAI_API_TOKEN:
|
43 |
+
headers["Authorization"] = f"Bearer {CIVITAI_API_TOKEN}"
|
44 |
+
resp = requests.get(url, headers=headers, stream=stream, timeout=60)
|
45 |
+
resp.raise_for_status()
|
46 |
+
return resp
|
47 |
+
|
48 |
+
def parse_civitai_model_id(url: str) -> Optional[str]:
|
49 |
+
m = re.search(r"civitai\.com/(?:models|api/v1/models)/(\d+)", url)
|
50 |
+
return m.group(1) if m else None
|
51 |
+
|
52 |
+
def fetch_civitai_model_json(url: str) -> Dict[str, Any]:
|
53 |
+
mid = parse_civitai_model_id(url)
|
54 |
+
if not mid:
|
55 |
+
raise gr.Error("Invalid CivitAI model URL.")
|
56 |
+
api_url = f"https://civitai.com/api/v1/models/{mid}"
|
57 |
+
return http_get(api_url).json()
|
58 |
+
|
59 |
+
def detect_nsfw(model_json: Dict[str, Any]) -> bool:
|
60 |
+
try:
|
61 |
+
if model_json.get("nsfw", False):
|
62 |
+
return True
|
63 |
+
if int(model_json.get("nsfwLevel") or 0) != 0:
|
64 |
+
return True
|
65 |
+
for mv in model_json.get("modelVersions", []):
|
66 |
+
if int(mv.get("nsfwLevel") or 0) != 0:
|
67 |
+
return True
|
68 |
+
for im in mv.get("images", []):
|
69 |
+
if int(im.get("nsfwLevel") or 0) != 0:
|
70 |
+
return True
|
71 |
+
return False
|
72 |
+
except Exception:
|
73 |
+
return True
|
74 |
+
|
75 |
+
def guess_base_model(model_json: Dict[str, Any]) -> Tuple[str, bool]:
|
76 |
+
base = ""
|
77 |
+
is_video = False
|
78 |
+
if "modelVersions" in model_json and model_json["modelVersions"]:
|
79 |
+
mv0 = model_json["modelVersions"][0]
|
80 |
+
base = (mv0.get("baseModel") or "").strip().lower()
|
81 |
+
is_video = "video" in base or "hunyuan" in base or "wan" in base
|
82 |
+
hf = BASE_MODEL_MAP.get(base, "stabilityai/stable-diffusion-xl-base-1.0")
|
83 |
+
return hf, is_video
|
84 |
+
|
85 |
+
def pick_primary_weight(model_json: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
86 |
+
candidates = []
|
87 |
+
for mv in model_json.get("modelVersions", []):
|
88 |
+
for f in mv.get("files", []):
|
89 |
+
url = f.get("downloadUrl") or f.get("url") or ""
|
90 |
+
if url.endswith(".safetensors"):
|
91 |
+
size = int(f.get("sizeKB", 0))
|
92 |
+
candidates.append((size, f))
|
93 |
+
if not candidates:
|
94 |
+
return None
|
95 |
+
candidates.sort(key=lambda x: x[0], reverse=True)
|
96 |
+
return candidates[0][1]
|
97 |
+
|
98 |
+
def download_file(url: str, dest: str):
|
99 |
+
r = http_get(url, stream=True)
|
100 |
+
with open(dest, "wb") as f:
|
101 |
+
for chunk in r.iter_content(chunk_size=1024 * 512):
|
102 |
+
if chunk: f.write(chunk)
|
103 |
+
|
104 |
+
def collect_previews(model_json: Dict[str, Any], dest_dir: str, max_items=6) -> List[str]:
|
105 |
+
paths = []
|
106 |
+
count = 0
|
107 |
+
for mv in model_json.get("modelVersions", []):
|
108 |
+
for im in mv.get("images", []):
|
109 |
+
if count >= max_items: break
|
110 |
+
url = im.get("url") or im.get("urlSmall") or ""
|
111 |
+
if not url: continue
|
112 |
+
ext = os.path.splitext(url.split("?")[0])[1] or ".jpg"
|
113 |
+
name = f"preview_{count+1}{ext}"
|
114 |
+
p = os.path.join(dest_dir, name)
|
115 |
+
try:
|
116 |
+
download_file(url, p)
|
117 |
+
paths.append(p); count += 1
|
118 |
+
except: continue
|
119 |
+
if count >= max_items: break
|
120 |
+
return paths
|
121 |
+
|
122 |
+
def extract_trained_words(model_json: Dict[str, Any]) -> List[str]:
|
123 |
+
words = []
|
124 |
+
for mv in model_json.get("modelVersions", []):
|
125 |
+
for w in (mv.get("trainedWords") or []):
|
126 |
+
if isinstance(w, str) and w.strip():
|
127 |
+
words.append(w.strip())
|
128 |
+
return list(dict.fromkeys(words))
|
129 |
+
|
130 |
+
def build_readme(out_dir: str, repo_name: str, model_json: Dict[str, Any],
|
131 |
+
hf_base_model: str, is_video: bool, is_nsfw: bool,
|
132 |
+
weight_filename: Optional[str], previews: List[str]) -> None:
|
133 |
+
name = model_json.get("name") or f"CivitAI-{model_json.get('id','model')}"
|
134 |
+
trained_words = extract_trained_words(model_json)
|
135 |
+
|
136 |
+
tags = ["lora","diffusers"]
|
137 |
+
if is_video: tags.append("video")
|
138 |
+
else: tags += ["text-to-image","stable-diffusion","template:sd-lora"]
|
139 |
+
if is_nsfw: tags.append("not-for-all-audiences")
|
140 |
+
|
141 |
+
yaml_header = f"""---
|
142 |
+
tags:
|
143 |
+
- {'\n - '.join(sorted(set(tags)))}
|
144 |
+
base_model: {hf_base_model}
|
145 |
+
instance_prompt: {trained_words[0] if trained_words else ''}
|
146 |
+
---
|
147 |
+
"""
|
148 |
+
|
149 |
+
diffusers_md = f"""
|
150 |
+
```py
|
151 |
+
from diffusers import AutoPipelineForText2Image
|
152 |
+
import torch
|
153 |
+
|
154 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
155 |
+
pipe = AutoPipelineForText2Image.from_pretrained("{hf_base_model}").to(device)
|
156 |
+
pipe.load_lora_weights("{repo_name}", weight_name="{weight_filename or 'your_lora.safetensors'}")
|
157 |
+
image = pipe("{trained_words[0] if trained_words else 'A sample prompt'}").images[0]
|
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
python-slugify
|
2 |
+
selenium
|
3 |
+
webdriver-manager
|
4 |
+
huggingface-hub==0.22.2
|
5 |
+
beautifulsoup4
|