Upload 3 files

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+RUN apt update && apt install -y \
+  rclone \
+  aria2 \
+  zip
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,529 @@
+import asyncio
+import datetime
+import json
+import logging
+import os
+import re
+import shutil
+import subprocess
+import time
+import uuid
+from typing import Optional
+
+import requests
+from bs4 import BeautifulSoup
+from fake_useragent import UserAgent
+from fastapi import FastAPI
+from huggingface_hub import HfApi, 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': UA.random,
+        "Content-Type": "application/json"
+    }
+
+
+class CivitAICrawler:
+    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス"""
+
+    def __init__(self, config: Config):
+        import base64
+
+        rclone_conf_base64 = os.environ.get("RCLONE_CONF_BASE64")
+        if rclone_conf_base64:
+            # カレントディレクトリ配下に .rclone_config ディレクトリを作成
+            config_dir = os.path.join(os.getcwd(), ".rclone_config")
+            os.makedirs(config_dir, exist_ok=True)
+
+            conf_path = os.path.join(config_dir, "rclone.conf")
+            with open(conf_path, "wb") as f:
+                f.write(base64.b64decode(rclone_conf_base64))
+
+            # rclone がここを参照するように設定
+            os.environ["RCLONE_CONFIG"] = conf_path
+            logger.info(f"[INFO] Created rclone.conf at {conf_path}")
+        else:
+            logger.warning("[WARN] RCLONE_CONF_BASE64 not found; rclone may fail.")
+        # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝
+
+        self.config = config
+        self.api = HfApi()
+        self.app = FastAPI()
+        self.repo_ids = self.config.REPO_IDS.copy()
+        self.jst = self.config.JST
+        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です。
+            モデル一覧は https://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.log を参照してください。
+            Status: {now} + currently running :D
+            """
+            return description
+
+        @self.app.on_event("startup")
+        async def startup_event():
+            asyncio.create_task(self.crawl())
+
+    @staticmethod
+    def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
+        """Content-Dispositionヘッダーからファイル名を取得する。"""
+        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) -> Optional[str]:
+        """指定されたURLからファイルをダウンロードし、指定されたフォルダに保存する。"""
+        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 None
+
+        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"Downloaded: {file_path}")
+        return 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}")
+            return {}
+
+    def download_model_files(self, model_versions: list, folder: str):
+        """最新のモデルバージョンと古いバージョンのファイルをまとめてダウンロード."""
+        for version in model_versions:
+            files_info = version.get("files", [])
+            for file_info in files_info:
+                download_url = file_info["downloadUrl"]
+                file_name = file_info["name"]
+                login_detected_count = 0
+
+                while login_detected_count < 5:
+                    local_path = self.download_file(download_url, folder, file_name)
+                    if local_path and "login" in os.listdir(folder):
+                        # 万が一、ダウンロード先に "login" という謎ファイルが出た場合の再試行処理
+                        login_detected_count += 1
+                        os.remove(os.path.join(folder, "login"))
+                        logger.warning(f"Detected 'login' file, retrying download: {file_name}  ({login_detected_count}/5)")
+                    else:
+                        break
+
+                if login_detected_count >= 5:
+                    # ダウンロード失敗を示すダミーファイルを作成
+                    dummy_file_path = os.path.join(folder, f"{file_name}.download_failed")
+                    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_path}")
+                    except Exception as e:
+                        logger.error(f"Failed to create dummy file for {file_name}: {e}")
+
+    def download_images(self, model_versions: list, folder: str):
+        """画像を images フォルダにまとめてダウンロードする."""
+        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 = os.path.basename(image_url)  # ファイル名部分
+            local_path = os.path.join(images_folder, image_name)
+            try:
+                resp = requests.get(image_url, stream=True)
+                resp.raise_for_status()
+                with open(local_path, 'wb') as imgf:
+                    for chunk in resp.iter_content(chunk_size=8192):
+                        imgf.write(chunk)
+                logger.info(f"Downloaded image: {local_path}")
+            except requests.RequestException as e:
+                logger.error(f"Failed to download image {image_url}: {e}")
+
+    def save_html_content(self, model_page_url: str, folder: str):
+        """モデルページのHTMLをフォルダ内に保存する."""
+        try:
+            resp = requests.get(model_page_url)
+            resp.raise_for_status()
+            html_path = os.path.join(folder, "page.html")
+            with open(html_path, 'w', encoding='utf-8') as f:
+                f.write(resp.text)
+            logger.info(f"Saved HTML: {html_path}")
+        except Exception as e:
+            logger.error(f"Error saving HTML content from {model_page_url}: {e}")
+
+    def save_model_info_json(self, model_info: dict, folder: str):
+        """モデル情報をJSONファイルとして保存."""
+        info_path = os.path.join(folder, "model_info.json")
+        try:
+            with open(info_path, 'w', encoding='utf-8') as f:
+                json.dump(model_info, f, indent=2)
+            logger.info(f"Saved model_info.json: {info_path}")
+        except Exception as e:
+            logger.error(f"Failed to save model info JSON: {e}")
+
+    def encrypt_and_upload_folder(self, local_folder: str):
+        """
+        1. rcloneでフォルダ全体を暗号化 (フォルダ名含む)
+        2. 暗号化されたフォルダをHugging Faceにアップロード
+        3. ローカル削除
+        """
+        if not os.path.exists(local_folder):
+            logger.error(f"encrypt_and_upload_folder: folder not found: {local_folder}")
+            return
+
+        # 暗号化後のフォルダが生成されるベースパス (例: /app/encrypted)
+        encrypted_base_dir = "/app/encrypted"
+        os.makedirs(encrypted_base_dir, exist_ok=True)
+
+        # rclone実行前の /app/encrypted の状態を取得 (新規フォルダ検出用)
+        before_set = set(os.listdir(encrypted_base_dir))
+
+        # rcloneでフォルダごとコピー (ファイル名・フォルダ名ともに暗号化)
+        #   ここで "cryptLocal:" は .rclone.conf 側で
+        #     [cryptLocal]
+        #     type = crypt
+        #     remote = /app/encrypted
+        #     filename_encryption = standard
+        #     password = ****
+        #   等が設定されている想定
+        try:
+            subprocess.run(
+                ["rclone", "copy", local_folder, "cryptLocal:"],
+                check=True
+            )
+        except subprocess.CalledProcessError as e:
+            logger.error(f"rclone copy failed: {e}")
+            return
+
+        # rclone実行後の /app/encrypted の状態
+        after_set = set(os.listdir(encrypted_base_dir))
+        # 新しくできた暗号化フォルダを特定
+        new_folders = after_set - before_set
+        if not new_folders:
+            logger.error("No new encrypted folder found. Something went wrong.")
+            return
+
+        # 通常は1個のはずだが、複数あるなら先頭だけ使う
+        enc_folder_name = list(new_folders)[0]
+        enc_folder_path = os.path.join(encrypted_base_dir, enc_folder_name)
+
+        # Hugging Face上で、この暗号化フォルダをそのままアップロード
+        #   => HF側もフォルダ名が暗号化された状態で表示されます
+        try:
+            # path_in_repo も同じ暗号化名を指定
+            self.upload_folder(enc_folder_path, path_in_repo=enc_folder_name)
+            logger.info(f"Uploaded encrypted folder to HF: {enc_folder_path}")
+        except Exception as e:
+            logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")
+
+        # ローカル削除（平文のフォルダ & 暗号化済みフォルダ）
+        try:
+            shutil.rmtree(local_folder)
+            shutil.rmtree(enc_folder_path)
+            logger.info(f"Removed local folder: {local_folder} and encrypted folder: {enc_folder_path}")
+        except Exception as e:
+            logger.error(f"Failed to remove local folders: {e}")
+
+    def upload_file(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"Uploaded file: {file_path} to {repo_id} at {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("File limit exceeded, creating a new repo.")
+                    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
+                    continue
+                elif "you can retry this action in about 1 hour" in error_message:
+                    logger.warning("Rate limit hit. Waiting 1 hour...")
+                    time.sleep(3600)
+                    attempt -= 1
+                else:
+                    if attempt < max_retries:
+                        logger.warning(f"Failed to upload {file_path}, retry {attempt}/{max_retries}")
+                    else:
+                        logger.error(f"Failed after {max_retries} attempts: {e}")
+                        raise
+
+    def upload_folder(self, folder_path: str, path_in_repo: Optional[str] = None):
+        """
+        フォルダを Hugging Face リポジトリに一括アップロード
+        """
+        if path_in_repo is None:
+            path_in_repo = os.path.basename(folder_path)
+
+        max_retries = 5
+        attempt = 0
+        while attempt < max_retries:
+            try:
+                self.api.upload_folder(
+                    folder_path=folder_path,
+                    repo_id=self.repo_ids['current'],
+                    path_in_repo=path_in_repo
+                )
+                logger.info(f"Uploaded folder: {folder_path} to {self.repo_ids['current']} at {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("File limit exceeded, creating a new repo.")
+                    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
+                    continue
+                elif "you can retry this action in about 1 hour" in error_message:
+                    logger.warning("Rate limit hit. Waiting 1 hour...")
+                    time.sleep(3600)
+                    attempt -= 1
+                else:
+                    if attempt < max_retries:
+                        logger.warning(f"Failed to upload folder {folder_path}, retry {attempt}/{max_retries}")
+                    else:
+                        logger.error(f"Failed after {max_retries} attempts: {e}")
+                        raise
+
+    @staticmethod
+    def increment_repo_name(repo_id: str) -> str:
+        """リポジトリ名の末尾の数字をインクリメントする。"""
+        match = re.search(r'(\d+)$', repo_id)
+        if match:
+            number = int(match.group(1)) + 1
+            new_repo_id = re.sub(r'\d+$', str(number), repo_id)
+        else:
+            new_repo_id = f"{repo_id}1"
+        return new_repo_id
+
+    def read_model_list(self) -> dict:
+        """モデルリストを読み込む。"""
+        model_list = {}
+        try:
+            with open(self.config.LIST_FILE, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        parts = line.split(": ", 1)
+                        if len(parts) == 2:
+                            modelpage_name, model_hf_url = parts
+                            model_list[model_hf_url] = modelpage_name
+        except Exception as e:
+            logger.error(f"Failed to read model list: {e}")
+        return model_list
+
+    def get_repo_info(self, repo_id):
+        """リポジトリの情報を取得する。"""
+        try:
+            repo_info = self.api.repo_info(repo_id=repo_id, files_metadata=True)
+            file_paths = [sibling.rfilename for sibling in repo_info.siblings]
+            return file_paths
+        except Exception as e:
+            logger.error(f"Failed to get repo info for {repo_id}: {e}")
+            return []
+
+    def process_model(self, model_url: str):
+        """1つのモデルをダウンロードしてフォルダ丸ごと暗号化＆アップロードする."""
+        try:
+            # model_idを取得
+            model_id = model_url.rstrip("/").split("/")[-1]
+
+            # モデル情報を取得
+            model_info = self.get_model_info(model_id)
+            if not model_info or "modelVersions" not in model_info:
+                logger.error(f"No valid model info for ID {model_id}. Skipping.")
+                return
+
+            # バージョン一覧
+            versions = model_info["modelVersions"]
+            if not versions:
+                logger.warning(f"No modelVersions found for ID {model_id}.")
+                return
+
+            # フォルダ名として適当な名前をつける
+            # たとえばモデル名をベースにフォルダを作る(被り防止にUUIDを付与)
+            folder_name = model_info.get("name", "UnknownModel")
+            folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)  # OSで使えない文字除去
+            folder_name += "_" + str(uuid.uuid4())[:8]
+            os.makedirs(folder_name, exist_ok=True)
+
+            # モデルファイルをダウンロード
+            self.download_model_files(versions, folder_name)
+
+            # 画像を images/ にダウンロード
+            self.download_images(versions, folder_name)
+
+            # HTMLを取得
+            model_page_url = f"{self.config.URLS['modelPage']}{model_id}"
+            self.save_html_content(model_page_url, folder_name)
+
+            # model_info.json保存
+            self.save_model_info_json(model_info, folder_name)
+
+            # ここでフォルダごと暗号化＆アップロード
+            self.encrypt_and_upload_folder(folder_name)
+
+            # model_list.logに追記 (暗号化フォルダを直接参照するURLは分からないため、
+            #  とりあえず元の modelPage名 とかモデルIDのメモを書くだけに留める)
+            hf_url_placeholder = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/[ENCRYPTED_FOLDER]"
+            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
+                f.write(f"{model_info.get('name', 'UnnamedModel')} (ID:{model_id}): {hf_url_placeholder}\n")
+
+        except Exception as e:
+            logger.error(f"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を最新化
+                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 ""
+
+                # 新着モデル確認
+                r = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
+                r.raise_for_status()
+                latest_models = r.json().get("items", [])
+                latest_model_ids = [m["id"] for m in latest_models if "id" in m]
+
+                new_models = list(set(latest_model_ids) - set(old_models))
+                if new_models:
+                    logger.info(f"New model IDs found: {new_models}")
+                    model_id = new_models[0]
+
+                    for attempt in range(1, 6):
+                        try:
+                            self.process_model(self.config.URLS["modelId"] + str(model_id))
+                            break
+                        except Exception as e:
+                            logger.error(f"Failed to process model {model_id} (attempt {attempt}/5): {e}")
+                            if attempt == 5:
+                                logger.error(f"Skipping model {model_id} after 5 failures.")
+                            else:
+                                await asyncio.sleep(2)
+
+                    # 新モデルを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 with new model ID: {model_id}")
+
+                    # ログファイル & model_list.logをアップロード
+                    self.upload_file(
+                        file_path=self.config.LOG_FILE,
+                        repo_id=self.repo_ids["log"],
+                        path_in_repo=self.config.LOG_FILE
+                    )
+                    self.upload_file(
+                        file_path=self.config.LIST_FILE,
+                        repo_id=self.repo_ids["model_list"],
+                        path_in_repo=self.config.LIST_FILE
+                    )
+                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"No new models. Updated log: {self.config.LOG_FILE}")
+                    self.upload_file(
+                        file_path=self.config.LOG_FILE,
+                        repo_id=self.repo_ids["log"],
+                        path_in_repo=self.config.LOG_FILE
+                    )
+                    logger.info("Uploaded log file.")
+                    await asyncio.sleep(60)
+                    continue
+
+            except Exception as e:
+                logger.error(f"Error in crawl loop: {e}")
+                await asyncio.sleep(300)
+
+
+# FastAPIアプリケーション
+config = Config()
+crawler = CivitAICrawler(config)
+app = crawler.app

requirements.txt

@@ -0,0 +1,9 @@
+fastapi==0.74.*
+sentencepiece==0.1.*
+torch==1.11.*
+transformers==4.*
+uvicorn[standard]==0.17.*
+requests==2.27.*
+beautifulsoup4
+huggingface_hub
+fake-useragent