AiAF commited on
Commit
a8380df
·
verified ·
1 Parent(s): 987b872

Upload folder using huggingface_hub

Browse files
Files changed (2) hide show
  1. app.py +157 -0
  2. 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