Upload main.py

main.py

@@ -11,12 +11,16 @@ 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, hf_hub_download, login
+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"]
@@ -36,6 +40,7 @@ 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': 'civitai-crawler/1.0',
@@ -49,7 +54,7 @@ class Config:
 
 
 class CivitAICrawler:
-    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス（フォルダ名も暗号化対応版）"""
+    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス"""
 
     def __init__(self, config: Config):
         self.config = config
@@ -64,12 +69,14 @@ class CivitAICrawler:
         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
+            model_list.log や civitai_backup.log は暗号化しないでアップロードします。
+            モデルのフォルダやファイルは暗号化してアップロードします。
             Status: {now} + currently running :D
             """
             return description
@@ -79,7 +86,7 @@ class CivitAICrawler:
             asyncio.create_task(self.crawl())
 
     # =============================================================================
-    # rclone まわり
+    # rclone の設定・暗号化アップロード処理
     # =============================================================================
     def setup_rclone_conf(self):
         """環境変数 RCLONE_CONF_BASE64 から rclone.conf を生成し、RCLONE_CONFIG 環境変数を設定"""
@@ -97,12 +104,8 @@ class CivitAICrawler:
 
     def encrypt_with_rclone(self, local_path: str):
         """
-        指定ファイル or ディレクトリを cryptLocal:にコピー。
-        このとき「トップレベルのフォルダ名」も暗号化して保持するため、
-        cryptLocal:{os.path.basename(local_path)} の形でコピーする。
-        
-        ※ rclone の crypt設定が filename_encryption = standard 等なら、
-          フォルダ名やファイル名も丸ごと暗号化される
+        指定ファイル or ディレクトリを cryptLocal: にコピー。
+        フォルダ構造やファイル名を rclone の filename_encryption 設定に応じて暗号化する。
         """
         if not os.path.exists(local_path):
             raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}")
@@ -111,13 +114,11 @@ class CivitAICrawler:
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
-        # コピー先を cryptLocal:xxx に指定することで、xxx というフォルダ単位で
-        # 暗号化されたディレクトリを作成できるようにする。
+        # コピー先: cryptLocal:{basename(local_path)}
         top_level_name = os.path.basename(local_path.rstrip("/"))
         if not top_level_name:
-            top_level_name = "unknown"
+            top_level_name = "unnamed"
 
-        # rclone copy local_path -> cryptLocal:top_level_name
         cmd = ["rclone", "copy", local_path, f"cryptLocal:{top_level_name}", "-v"]
         logger.info(f"[INFO] Running: {' '.join(cmd)}")
         subprocess.run(cmd, check=True)
@@ -130,24 +131,16 @@ class CivitAICrawler:
 
     def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
         """
-        self.config.ENCRYPTED_DIR にある暗号化後のファイル/フォルダ構造を
-        そのまま Hugging Face にアップロードする。
-        
-        rclone で filename_encryption=standard が有効な場合:
-         - 元フォルダ名/ファイル名は完全に暗号化され、HF上では判読不能な名称になる。
-        
-        base_path_in_repo が空("")の場合は、
-        `/home/user/app/encrypted` の構造が HF リポジトリの直下に展開される。
+        self.config.ENCRYPTED_DIR 以下に生成された暗号化後のファイル・フォルダ構造を再帰的��
+        Hugging Face にアップロードする。フォルダ名・ファイル名は rclone により暗号化済み。
         """
         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
 
-                # self.config.ENCRYPTED_DIR からの相対パス（暗号化後のフォルダ名・ファイル名）
                 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)
 
@@ -372,28 +365,103 @@ class CivitAICrawler:
         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
+    ):
+        """
+        暗号化せず、そのまま Hugging Face にアップロードするメソッド。
+        civitai_backup.log や model_list.log などはこれを使う。
+        """
+        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 リセット
+                    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(self, file_path: str):
+    def upload_file_encrypted(
+        self,
+        file_path: str,
+        repo_id: Optional[str] = None,
+        path_in_repo: Optional[str] = None
+    ):
         """
-        単一ファイルを暗号化 → 暗号化後のものをHugging Faceへアップロード。
-        rcloneでフォルダを作って暗号化するため、フォルダ構造としてアップされる。
+        単一ファイルを rclone で暗号化し、そのままHFへアップロード。
         """
-        self.encrypt_with_rclone(file_path)  # 1) ローカルを -> cryptLocal:basename(…)
-        self.upload_encrypted_files(repo_id=self.repo_ids['current'], base_path_in_repo="")
+        if repo_id is None:
+            repo_id = self.repo_ids['current']
+        base_path = path_in_repo or ""  # HF上のベースパス
+
+        # 1) rclone で暗号化
+        self.encrypt_with_rclone(file_path)
+
+        # 2) アップロード
+        self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=base_path)
+
+        # 3) 後始末
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
-    def upload_folder(self, folder_path: str):
+    # =============================================================================
+    # ★ 暗号化してアップロード (フォルダ)
+    # =============================================================================
+    def upload_folder_encrypted(
+        self,
+        folder_path: str,
+        repo_id: Optional[str] = None,
+        path_in_repo: Optional[str] = None
+    ):
         """
-        フォルダを暗号化 → フォルダ名ごとアップロード。
-        rclone copy folder_path => cryptLocal:folder_pathのbasename
-        すると {ENCRYPTED_DIR}/{暗号化後のトップレベルフォルダ} が生成されるため、
-        それを再帰的にHugging Faceへアップロードする。
+        フォルダを rclone で暗号化し、暗号化されたフォルダ構造ごとアップロード。
         """
+        if repo_id is None:
+            repo_id = self.repo_ids['current']
+        base_path = path_in_repo or ""
+
         self.encrypt_with_rclone(folder_path)
-        self.upload_encrypted_files(repo_id=self.repo_ids['current'], base_path_in_repo="")
+        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)
 
@@ -462,7 +530,7 @@ class CivitAICrawler:
             self.save_model_info(model_info, folder)
 
             # ====== rclone でフォルダ暗号化 → HFへアップロード ======
-            self.upload_folder(folder)
+            self.upload_folder_encrypted(folder)
 
             # モデルリスト更新
             modelpage_name = model_info.get("name", "Unnamed Model")
@@ -527,12 +595,12 @@ class CivitAICrawler:
                     logger.info(f"Updated log file: {self.config.LOG_FILE}")
 
                     # ログファイルをリポジトリにアップロード
-                    self.upload_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.")
+                    logger.info("Uploaded log file to repository (unencrypted).")
 
                     logger.info("No new models found.")
                     await asyncio.sleep(60)
@@ -547,13 +615,13 @@ class CivitAICrawler:
                     f.write(f"{self.repo_ids['current']}\n")
                 logger.info(f"Updated log file with new model ID: {model_id}")
 
-                # ログとモデルリストのアップロード (rcloneは使わず直接uploadでもOKならそのまま)
-                self.upload_file(
+                # ログとモデルリストのアップロード
+                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(
+                self.upload_file_raw(
                     file_path=self.config.LIST_FILE,
                     repo_id=self.repo_ids["model_list"],
                     path_in_repo=self.config.LIST_FILE