import asyncio import base64 import datetime import json import logging import os import re import shutil import subprocess import time from typing import Optional import requests from bs4 import BeautifulSoup from fake_useragent import UserAgent from fastapi import FastAPI from huggingface_hub import HfApi, create_repo, hf_hub_download, login # ロギングの設定 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class Config: """設定用のクラス""" HUGGINGFACE_API_KEY = os.environ["HUGGINGFACE_API_KEY"] CIVITAI_API_TOKEN = os.environ["CIVITAI_API_TOKEN"] LOG_FILE = "civitai_backup.log" LIST_FILE = "model_list.log" REPO_IDS = { "log": "ttttdiva/CivitAI_log_test", "model_list": "ttttdiva/CivitAI_model_info_test", "current": "" } URLS = { "latest": "https://civitai.com/api/v1/models?sort=Newest", "modelPage": "https://civitai.com/models/", "modelId": "https://civitai.com/api/v1/models/", "modelVersionId": "https://civitai.com/api/v1/model-versions/", "hash": "https://civitai.com/api/v1/model-versions/by-hash/" } JST = datetime.timezone(datetime.timedelta(hours=9)) UA = UserAgent() HEADERS = { 'Authorization': f'Bearer {CIVITAI_API_TOKEN}', 'User-Agent': 'civitai-crawler/1.0', "Content-Type": "application/json" } # ===== rclone 用の追加設定 ===== RCLONE_CONF_BASE64 = os.environ.get("RCLONE_CONF_BASE64", "") # 暗号化されたファイルが出力されるローカルディレクトリ(cryptLocal: の実体) ENCRYPTED_DIR = "/home/user/app/encrypted" class CivitAICrawler: """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス""" def __init__(self, config: Config): self.config = config self.api = HfApi() self.app = FastAPI() self.repo_ids = self.config.REPO_IDS.copy() self.jst = self.config.JST # rclone のセットアップ self.setup_rclone_conf() self.setup_routes() def setup_routes(self): """FastAPIのルーティングを設定する。""" @self.app.get("/") def read_root(): now = str(datetime.datetime.now(self.jst)) description = f""" CivitAIを定期的に周回し新規モデルを {self.repo_ids['current']} にバックアップするSpaceです。 model_list.log や civitai_backup.log は暗号化しないでアップロードします。 モデルのフォルダやファイルは暗号化してアップロードします。 Status: {now} + currently running :D """ return description @self.app.on_event("startup") async def startup_event(): asyncio.create_task(self.crawl()) # ============================================================================= # rclone の設定・暗号化アップロード処理 # ============================================================================= def setup_rclone_conf(self): """環境変数 RCLONE_CONF_BASE64 から rclone.conf を生成し、RCLONE_CONFIG 環境変数を設定""" if not self.config.RCLONE_CONF_BASE64: logger.warning("[WARN] RCLONE_CONF_BASE64 is empty. rclone may fail.") return os.makedirs(".rclone_config", exist_ok=True) conf_path = os.path.join(".rclone_config", "rclone.conf") with open(conf_path, "wb") as f: f.write(base64.b64decode(self.config.RCLONE_CONF_BASE64)) os.environ["RCLONE_CONFIG"] = conf_path logger.info(f"[INFO] rclone.conf created at: {conf_path}") def encrypt_with_rclone(self, local_path: str): """ local_path がファイルかフォルダかで処理を分け、 - フォルダ → 従来通り rclone copy - ファイル → rclone copyto を使って暗号化 """ if not os.path.exists(local_path): raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}") # 事前に暗号先ディレクトリをクリーンアップ if os.path.isdir(self.config.ENCRYPTED_DIR): shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True) # ========= 1) local_path がフォルダの場合 ========= if os.path.isdir(local_path): top_level_name = os.path.basename(local_path.rstrip("/")) or "unnamed" cmd = [ "rclone", "copy", local_path, f"cryptLocal:{top_level_name}", "-v" ] logger.info(f"[INFO] Running: {' '.join(cmd)}") subprocess.run(cmd, check=True) # ========= 2) local_path がファイルの場合 ============= else: filename = os.path.basename(local_path) # "rclone copyto cryptLocal:" にすると # cryptLocal 側では1ファイルとして扱われる cmd = [ "rclone", "copyto", local_path, f"cryptLocal:{filename}", "-v" ] logger.info(f"[INFO] Running: {' '.join(cmd)}") subprocess.run(cmd, check=True) logger.info(f"[OK] rclone => cryptLocal:") if not os.path.isdir(self.config.ENCRYPTED_DIR): raise FileNotFoundError( f"[ERROR] {self.config.ENCRYPTED_DIR} not found. Check rclone config." ) # 例: upload_encrypted_files の中の再試行処理 def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""): max_retries = 5 for root, dirs, files in os.walk(self.config.ENCRYPTED_DIR): for fn in files: encrypted_file_path = os.path.join(root, fn) if not os.path.isfile(encrypted_file_path): continue relative_path = os.path.relpath(encrypted_file_path, self.config.ENCRYPTED_DIR) upload_path_in_repo = os.path.join(base_path_in_repo, relative_path) attempt = 0 while attempt < max_retries: try: self.api.upload_file( path_or_fileobj=encrypted_file_path, repo_id=repo_id, path_in_repo=upload_path_in_repo ) logger.info(f"[OK] Uploaded => {repo_id}/{upload_path_in_repo}") break except Exception as e: attempt += 1 error_message = str(e) # ================================ # 429によるrate-limit検出追加 # ================================ # "You have been rate-limited; you can retry this action in 31 minutes." # のようなメッセージから時間を抽出し、その時間+1分だけ待機後、再試行 if "rate-limited" in error_message and "minutes" in error_message: import re match = re.search(r"in (\d+) minutes?", error_message) if match: minutes = int(match.group(1)) # +1分して待機 minutes += 1 logger.warning(f"Rate-limited. Waiting {minutes} minutes before retry...") time.sleep(minutes * 60) attempt -= 1 # 同じ attempt カウントで再試行 continue # ================================ # すでにある1時間待機処理 # ================================ if "you can retry this action in about 1 hour" in error_message: logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour before retrying...") time.sleep(3600) attempt -= 1 # 再試行回数を増やさずにループを続ける continue if "over the limit of 100000 files" in error_message: logger.warning("Repository file limit exceeded. Creating a new repository...") self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current']) self.api.create_repo(repo_id=self.repo_ids['current'], private=True) attempt = 0 repo_id = self.repo_ids['current'] continue # 上記以外のエラーの場合 if attempt < max_retries: logger.warning( f"Failed to upload {encrypted_file_path}, retry {attempt}/{max_retries}..." ) else: logger.error( f"Failed to upload after {max_retries} attempts: {encrypted_file_path}" ) raise @staticmethod def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str: if content_disposition: parts = content_disposition.split(';') for part in parts: if "filename=" in part: return part.split("=")[1].strip().strip('"') return default_name def download_file(self, url: str, destination_folder: str, default_name: str): try: response = requests.get(url, headers=self.config.HEADERS, stream=True) response.raise_for_status() except requests.RequestException as e: logger.error(f"Failed to download file from {url}: {e}") return filename = self.get_filename_from_cd(response.headers.get('content-disposition'), default_name) file_path = os.path.join(destination_folder, filename) with open(file_path, 'wb') as file: for chunk in response.iter_content(chunk_size=8192): file.write(chunk) logger.info(f"Download completed: {file_path}") def get_model_info(self, model_id: str) -> dict: try: response = requests.get(self.config.URLS["modelId"] + str(model_id), headers=self.config.HEADERS) response.raise_for_status() return response.json() except requests.RequestException as e: logger.error(f"Failed to retrieve model info for ID {model_id}: {e}") def download_model(self, model_versions: list, folder: str): """ - 最新バージョンのファイルを「folder」へまとめてダウンロード - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード) (※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに) """ latest_version = model_versions[0] latest_files = latest_version["files"] for file_info in latest_files: download_url = file_info["downloadUrl"] file_name = file_info["name"] login_detected_count = 0 while login_detected_count < 5: try: self.download_file(download_url, folder, file_name) except Exception as e: logger.error(f"Exception occurred while downloading {file_name}: {e}") login_detected_count += 1 continue if "login" in os.listdir(folder): login_detected_count += 1 logger.warning(f"'login' file found. Will try again. ({login_detected_count}/5)") os.remove(os.path.join(folder, "login")) else: logger.info(f"Successfully downloaded {file_name}") break if login_detected_count >= 5: dummy_file_name = f"{file_name}.download_failed" dummy_file_path = os.path.join(folder, dummy_file_name) try: with open(dummy_file_path, "w") as f: f.write("Download failed after 5 attempts.") logger.error(f"Failed to download {file_name}. Created dummy file {dummy_file_name}. URL: {download_url}") except Exception as e: logger.error(f"Failed to create dummy file for {file_name}: {e}") def download_old_version_file_and_upload(self, file_info, parent_folder: str, encrypted_top_name: str): """ 古いバージョンのファイルを1つダウンロード→暗号化アップロード→削除。 「old_versions」というフォルダ名をそのまま使う。 """ file_name = file_info["name"] download_url = file_info["downloadUrl"] # old_versions フォルダを作成 old_versions_folder = os.path.join(parent_folder, "old_versions") os.makedirs(old_versions_folder, exist_ok=True) local_path = os.path.join(old_versions_folder, file_name) login_detected_count = 0 while login_detected_count < 5: try: # old_versionsフォルダへダウンロード self.download_file(download_url, old_versions_folder, file_name) except Exception as e: logger.error(f"Exception while downloading old file {file_name}: {e}") login_detected_count += 1 continue # 万が一 "login" という名前のファイルが出現したら削除して再試行 if "login" in os.listdir(old_versions_folder): login_detected_count += 1 logger.warning(f"'login' file found while downloading {file_name}. Retry {login_detected_count}/5.") os.remove(os.path.join(old_versions_folder, "login")) else: logger.info(f"Successfully downloaded old version file: {file_name}") break if login_detected_count >= 5: # 5回失敗 dummy_file_name = f"{file_name}.download_failed" dummy_file_path = os.path.join(old_versions_folder, dummy_file_name) try: with open(dummy_file_path, "w") as f: f.write("Download failed after 5 attempts.") logger.error(f"Failed to download {file_name} -> created dummy: {dummy_file_name}") except Exception as e: logger.error(f"Failed to create dummy file for old version {file_name}: {e}") return # ダウンロード成功 → 1ファイルだけ暗号化&アップロード → ローカル削除 try: # 1) ファイル単位で暗号化するなら self.encrypt_with_rclone(local_path) # 2) 暗号化ファイルをアップロード self.upload_encrypted_files( repo_id=self.repo_ids["current"], base_path_in_repo=f"{encrypted_top_name}/old_versions" ) logger.info(f"Uploaded old version file: {file_name} into {encrypted_top_name}") except Exception as e: logger.error(f"Error uploading old version file {file_name}: {e}") # 3) アップロード後、平文ファイルを削除 if os.path.exists(local_path): os.remove(local_path) logger.info(f"Removed local old version file: {local_path}") def download_images(self, model_versions: list, folder: str): images_folder = os.path.join(folder, "images") os.makedirs(images_folder, exist_ok=True) images = [] for version in model_versions: for img in version.get("images", []): image_url = img["url"] images.append(image_url) for image_url in images: image_name = image_url.split("/")[-1] try: response = requests.get(image_url) response.raise_for_status() with open(os.path.join(images_folder, f"{image_name}.png"), "wb") as file: file.write(response.content) except requests.RequestException as e: logger.error(f"Error downloading image {image_url}: {e}") def save_html_content(self, url: str, folder: str): try: response = requests.get(url) response.raise_for_status() html_path = os.path.join(folder, f"{folder}.html") with open(html_path, 'w', encoding='utf-8') as file: file.write(response.text) except Exception as e: logger.error(f"Error saving HTML content for URL {url}: {e}") @staticmethod def save_model_info(model_info: dict, folder: str): with open(os.path.join(folder, "model_info.json"), "w") as file: json.dump(model_info, file, indent=2) @staticmethod def increment_repo_name(repo_id: str) -> str: match = re.search(r'(\d+)$', repo_id) if match: number = int(match.group(1)) + 1 return re.sub(r'\d+$', str(number), repo_id) else: return f"{repo_id}1" # ============================================================================= # 暗号化しないアップロード(ログや model_list.log 用) # ============================================================================= def upload_file_raw( self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None ): if repo_id is None: repo_id = self.repo_ids['current'] if path_in_repo is None: path_in_repo = os.path.basename(file_path) max_retries = 5 attempt = 0 while attempt < max_retries: try: self.api.upload_file( path_or_fileobj=file_path, repo_id=repo_id, path_in_repo=path_in_repo ) logger.info(f"[OK] Uploaded {file_path} => {repo_id}/{path_in_repo}") return except Exception as e: attempt += 1 error_message = str(e) if "over the limit of 100000 files" in error_message: logger.warning("Repository file limit exceeded, creating a new repository.") self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current']) self.api.create_repo(repo_id=self.repo_ids['current'], private=True) attempt = 0 repo_id = self.repo_ids['current'] continue elif "you can retry this action in about 1 hour" in error_message: logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour before retrying...") time.sleep(3600) attempt -= 1 else: if attempt < max_retries: logger.warning(f"Failed to upload raw file {file_path}, retry {attempt}/{max_retries}...") else: logger.error(f"Failed to upload raw file after {max_retries} attempts: {file_path}") raise # ============================================================================= # 暗号化してアップロード (単ファイル) # ============================================================================= def upload_file_encrypted( self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None ): if repo_id is None: repo_id = self.repo_ids['current'] base_path = path_in_repo or "" self.encrypt_with_rclone(file_path) self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=base_path) if os.path.isdir(self.config.ENCRYPTED_DIR): shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True) # ============================================================================= # 暗号化してアップロード (フォルダ) # ============================================================================= def upload_folder_encrypted( self, folder_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None ) -> str: if repo_id is None: repo_id = self.repo_ids['current'] base_path = path_in_repo or "" self.encrypt_with_rclone(folder_path) top_levels = [ d for d in os.listdir(self.config.ENCRYPTED_DIR) if os.path.isdir(os.path.join(self.config.ENCRYPTED_DIR, d)) ] if not top_levels: raise RuntimeError("No top-level folder found after rclone encryption.") if len(top_levels) > 1: logger.warning(f"Multiple top-level folders found after encryption? {top_levels}. Using the first one.") encrypted_top_name = top_levels[0] self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=base_path) if os.path.isdir(self.config.ENCRYPTED_DIR): shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True) return encrypted_top_name # ============================================================================= # model_list.log の読み書きを「model_id: model_hf_url」で扱うよう変更 # ============================================================================= def read_model_list(self): """ model_list.log の各行を "123456: https://huggingface.co/...encrypted_folder_name" の形式で読み込み、 { "123456": "https://huggingface.co/..."} の dict を返す """ model_list = {} try: with open(self.config.LIST_FILE, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue parts = line.split(": ", 1) if len(parts) == 2: stored_id, stored_url = parts model_list[stored_id] = stored_url return model_list except Exception as e: logger.error(f"Failed to read model list: {e}") return {} def process_model(self, model_url: str): """ - 最新バージョン(+ images, html, info.json)をフォルダに一括DL → 最後にまとめて暗号化アップロード - old_versions のファイルは 1つずつ「old_versions/」へDL → 即暗号化アップロード """ try: model_id = model_url.rstrip("/").split("/")[-1] model_info = self.get_model_info(model_id) model_versions = model_info.get("modelVersions", []) if not model_versions: logger.warning(f"No versions found for model ID {model_id}") return # ================================================================== # 1) 「古いコード」と同じロジックでフォルダ名を決定 # ================================================================== latest_version = model_versions[0] model_file = next( (file for file in latest_version["files"] if file.get('type') == 'Model'), None ) if model_file: latest_filename = model_file['name'] folder = os.path.splitext(latest_filename)[0] else: first_file = latest_version["files"][0] latest_filename = first_file['name'] folder = os.path.splitext(latest_filename)[0] logger.warning(f"No 'Model' type file found for model ID {model_id}. Using first file's name.") os.makedirs(folder, exist_ok=True) # ================================================================== # 2) 最新バージョンをまとめてダウンロード (model本体) # + 画像ダウンロード + HTML保存 + info.json保存 # ================================================================== self.download_model(model_versions, folder) self.download_images(model_versions, folder) self.save_html_content(model_url, folder) self.save_model_info(model_info, folder) # ================================================================== # 3) フォルダ全体(最新バージョン + images + html + info.json)を # まとめて 1 回だけ暗号化→アップロード # ================================================================== encrypted_top_name = self.upload_folder_encrypted(folder) logger.info(f"[MAIN] Uploaded latest version folder => {encrypted_top_name}") # ローカルフォルダ削除 shutil.rmtree(folder, ignore_errors=True) # ================================================================== # 4) 古いバージョンのファイルを 1つずつDL → 即アップロード # ================================================================== if len(model_versions) > 1: # 旧バージョンのファイルだけ個別処理 for version in model_versions[1:]: for file_info in version["files"]: # ***** 修正ポイント ***** # download_old_version_file_and_upload(file_info, parent_folder, encrypted_top_name) self.download_old_version_file_and_upload(file_info, folder, encrypted_top_name) # ↑ parent_folder=folder に old_versionsフォルダを作り # そこへダウンロードしてから暗号化アップロード # ================================================================== # 5) model_list.log への登録やその他の処理 # ================================================================== modelpage_name = model_info.get("name", f"Model_{model_id}") model_hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{encrypted_top_name}" with open(self.config.LIST_FILE, "a", encoding="utf-8") as f: f.write(f"{modelpage_name}: {model_hf_url}\n") except Exception as e: logger.error(f"Unexpected error in process_model ({model_url}): {e}") async def crawl(self): """モデルを定期的にチェックし、更新を行う。""" while True: try: login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True) # model_list.log & civitai_backup.log を取得 model_list_path = hf_hub_download(repo_id=self.repo_ids['model_list'], filename=self.config.LIST_FILE) shutil.copyfile(model_list_path, f"./{self.config.LIST_FILE}") local_file_path = hf_hub_download(repo_id=self.repo_ids["log"], filename=self.config.LOG_FILE) shutil.copyfile(local_file_path, f"./{self.config.LOG_FILE}") # ログ読み込み with open(self.config.LOG_FILE, "r", encoding="utf-8") as file: lines = file.read().splitlines() old_models = json.loads(lines[0]) if len(lines) > 0 else [] self.repo_ids["current"] = lines[1] if len(lines) > 1 else "" # 新着モデル確認 response = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS) response.raise_for_status() latest_models = response.json().get("items", []) latest_model_ids = [item.get("id") for item in latest_models if "id" in item] # 増分チェック new_models = list(set(latest_model_ids) - set(old_models)) if new_models: logger.info(f"New models found: {new_models}") model_id = new_models[0] for attempt in range(1, 6): try: self.process_model(f"{self.config.URLS['modelId']}{model_id}") break except Exception as e: logger.error(f"Failed to process model ID {model_id} (Attempt {attempt}/5): {e}") if attempt == 5: logger.error(f"Skipping model ID {model_id} after 5 failed attempts.") else: await asyncio.sleep(2) else: # 新モデルなし with open(self.config.LOG_FILE, "w", encoding="utf-8") as f: f.write(json.dumps(latest_model_ids) + "\n") f.write(f"{self.repo_ids['current']}\n") logger.info(f"Updated log file: {self.config.LOG_FILE}") self.upload_file_raw( file_path=self.config.LOG_FILE, repo_id=self.repo_ids["log"], path_in_repo=self.config.LOG_FILE ) logger.info("Uploaded log file to repository (unencrypted).") logger.info("No new models found.") await asyncio.sleep(60) continue # 追加したモデルIDを old_models に追加 old_models.append(model_id) # ログファイル更新 with open(self.config.LOG_FILE, "w", encoding="utf-8") as f: f.write(json.dumps(old_models) + "\n") f.write(f"{self.repo_ids['current']}\n") logger.info(f"Updated log file with new model ID: {model_id}") # ログとmodel_list.logをアップロード self.upload_file_raw( file_path=self.config.LOG_FILE, repo_id=self.repo_ids["log"], path_in_repo=self.config.LOG_FILE ) self.upload_file_raw( file_path=self.config.LIST_FILE, repo_id=self.repo_ids["model_list"], path_in_repo=self.config.LIST_FILE ) except Exception as e: logger.error(f"Error during crawling: {e}") await asyncio.sleep(300) # 実行 config = Config() crawler = CivitAICrawler(config) app = crawler.app