Upload main.py

main.py

@@ -11,16 +11,13 @@ 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
+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"]
@@ -40,57 +37,21 @@ class Config:
         "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,
+        'User-Agent': 'civitai-crawler/1.0',
         "Content-Type": "application/json"
     }
 
     # ===== rclone 用の追加設定 =====
-    # （環境変数 RCLONE_CONF_BASE64 に rclone.conf をbase64エンコードした文字列を設定しておく想定）
     RCLONE_CONF_BASE64 = os.environ.get("RCLONE_CONF_BASE64", "")
     # 暗号化されたファイルが出力されるローカルディレクトリ（cryptLocal: の実体）
     ENCRYPTED_DIR = "/home/user/app/encrypted"
 
 
-class CivitAICrawler:
-    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス（rcloneで暗号化対応版）"""
-
-    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です。
-            モデルページ名とバックアップURLの紐づきはhttps://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.logからどうぞ
-            たまに覗いてもらえると動き続けると思います。
-            再起動が必要になっている場合はRestartボタンを押してもらえると助かります。
-            Status: {now} + currently running :D
-            """
-            return description
-
-        @self.app.on_event("startup")
-        async def startup_event():
-            asyncio.create_task(self.crawl())
-
-    # =============================
-    # rclone 周りのヘルパー関数
-    # =============================
+    # =============================================================================
+    # rclone まわり
+    # =============================================================================
     def setup_rclone_conf(self):
         """環境変数 RCLONE_CONF_BASE64 から rclone.conf を生成し、RCLONE_CONFIG 環境変数を設定"""
         if not self.config.RCLONE_CONF_BASE64:
@@ -105,52 +66,49 @@ class CivitAICrawler:
         os.environ["RCLONE_CONFIG"] = conf_path
         logger.info(f"[INFO] rclone.conf created at: {conf_path}")
 
-    def encrypt_with_rclone(self, local_path: str, is_file: bool = True):
+    def encrypt_with_rclone(self, local_path: str):
         """
-        指定ファイル or フォルダを cryptLocal: へコピーし、暗号化ファイルを self.config.ENCRYPTED_DIR に生成する。
-        rclone copy の引数にフォルダパスやファイルパスを指定して利用可能。
+        指定ファイル or ディレクトリを cryptLocal: にコピーし、
+        暗号化された結果を self.config.ENCRYPTED_DIR に生成する。
         """
         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)
 
-        # rclone コマンドの実行
         cmd = ["rclone", "copy", local_path, "cryptLocal:", "-v"]
         logger.info(f"[INFO] Running: {' '.join(cmd)}")
         subprocess.run(cmd, check=True)
-        logger.info(f"[OK] rclone copy {local_path} => cryptLocal:")
+        logger.info(f"[OK] rclone copy => cryptLocal:")
 
         if not os.path.isdir(self.config.ENCRYPTED_DIR):
             raise FileNotFoundError(
                 f"[ERROR] {self.config.ENCRYPTED_DIR} not found. Check your rclone config."
             )
 
-    def upload_encrypted_files(self, repo_id: str, path_in_repo: str = None):
+    def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
         """
-        self.config.ENCRYPTED_DIR にある暗号化済みファイルをまとめて Hugging Face にアップロードする。
-        もとがフォルダの場合はサブディレクトリも含めて再帰的にアップロードする。
+        self.config.ENCRYPTED_DIR にある暗号化後のファイル/フォルダ構造をそっくりそのまま
+        Hugging Faceにアップロードする。フォルダ名・ファイル名はrcloneの設定によって暗号化されている。
         """
-        if not path_in_repo:
-            path_in_repo = ""
-
         max_retries = 5
 
-        # 再帰的に walk
+        # 再帰的に暗号化フォルダを探索
         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
 
-                # 元ディレクトリ相対のパスを生成
+                # self.config.ENCRYPTED_DIR からの相対パスを算出
                 relative_path = os.path.relpath(encrypted_file_path, self.config.ENCRYPTED_DIR)
-                # Hugging Face 上に保存するときのフルパス
-                upload_path_in_repo = os.path.join(path_in_repo, relative_path)
+                # Hugging Face上にアップロードする際のディレクトリパス
+                # ここでは base_path_in_repo + relative_path にしているが、
+                # base_path_in_repo が空("")なら、rclone で暗号化されたディレクトリ名をそのまま使う
+                upload_path_in_repo = os.path.join(base_path_in_repo, relative_path)
 
-                # HFへのアップロードを試行 (over the limitなどの例外をリトライする)
                 attempt = 0
                 while attempt < max_retries:
                     try:
@@ -159,16 +117,16 @@ class CivitAICrawler:
                             repo_id=repo_id,
                             path_in_repo=upload_path_in_repo
                         )
-                        logger.info(f"[OK] Uploaded {encrypted_file_path} to {repo_id}/{upload_path_in_repo}")
-                        break  # 成功したらループ抜け
+                        logger.info(f"[OK] Uploaded => {repo_id}/{upload_path_in_repo}")
+                        break
                     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.")
+                            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 をリセット
+                            # リポジトリが変わったので attempt をリセット
                             attempt = 0
                             repo_id = self.repo_ids['current']
                             continue
@@ -178,9 +136,13 @@ class CivitAICrawler:
                             attempt -= 1  # この場合はリトライ回数をカウントしない
                         else:
                             if attempt < max_retries:
-                                logger.warning(f"Failed to upload file {encrypted_file_path}, retrying... {attempt}/{max_retries}")
+                                logger.warning(
+                                    f"Failed to upload {encrypted_file_path}, retry {attempt}/{max_retries}..."
+                                )
                             else:
-                                logger.error(f"Failed to upload file after {max_retries} attempts: {encrypted_file_path}")
+                                logger.error(
+                                    f"Failed to upload after {max_retries} attempts: {encrypted_file_path}"
+                                )
                                 raise
 
     # =============================
@@ -367,44 +329,30 @@ class CivitAICrawler:
             new_repo_id = f"{repo_id}1"
         return new_repo_id
 
-    # =============================
-    # ここを rclone 暗号化＆アップロードに書き換え
-    # =============================
-    def upload_file(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None):
+    # =============================================================================
+    # 既存のファイル／フォルダアップロードを「フォルダ名も暗号化」するよう変更
+    # =============================================================================
+    def upload_file(self, file_path: str):
         """
         1) rcloneで file_path を暗号化
-        2) 暗号化されたファイル(群)を Hugging Face にアップロード
+        2) 暗号化された単一ファイル(またはフォルダ構造)をそのままHugging Faceへアップロード
+           （ファイル名もフォルダ名も暗号化される）
         """
-        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)
-
-        # 1) rclone copy (ファイル暗号化)
-        self.encrypt_with_rclone(file_path, is_file=True)
-
-        # 2) 暗号ファイルをアップロード
-        self.upload_encrypted_files(repo_id=repo_id, path_in_repo=path_in_repo)
-
-        # 3) 暗号ディレクトリの掃除
+        # rcloneで暗号化
+        self.encrypt_with_rclone(file_path)
+        # 暗号化フォルダ全体をアップロード (フォルダ名も暗号化される)
+        self.upload_encrypted_files(repo_id=self.repo_ids['current'], base_path_in_repo="")
+        # 後始末
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
-    def upload_folder(self, folder_path: str, path_in_repo: Optional[str] = None):
+    def upload_folder(self, folder_path: str):
         """
-        1) rcloneで folder_path を暗号化
-        2) 暗号化されたフォルダを Hugging Face にアップロード
+        1) rcloneでフォルダ全体を暗号化 (フォルダ名含むディレクトリ構造が暗号化)
+        2) 暗号化済みディレクトリ構造をHugging Faceにアップロード
         """
-        if path_in_repo is None:
-            path_in_repo = os.path.basename(folder_path)
-
-        # 1) rclone copy (フォルダ暗号化)
-        self.encrypt_with_rclone(folder_path, is_file=False)
-
-        # 2) 暗号フォルダをアップロード
-        self.upload_encrypted_files(repo_id=self.repo_ids['current'], path_in_repo=path_in_repo)
-
-        # 3) 掃除
+        self.encrypt_with_rclone(folder_path)
+        self.upload_encrypted_files(repo_id=self.repo_ids['current'], base_path_in_repo="")
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)