Upload main.py

main.py

@@ -14,7 +14,6 @@ import requests
 from fastapi import FastAPI
 from huggingface_hub import HfApi, hf_hub_download, login
 
-# ロギングの設定
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -48,6 +47,7 @@ class Config:
     # 暗号化されたファイルが出力されるローカルディレクトリ（cryptLocal: の実体）
     ENCRYPTED_DIR = "/home/user/app/encrypted"
 
+
 class CivitAICrawler:
     """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス（フォルダ名も暗号化対応版）"""
 
@@ -58,7 +58,7 @@ class CivitAICrawler:
         self.repo_ids = self.config.REPO_IDS.copy()
         self.jst = self.config.JST
 
-        # rclone のセットアップ（rclone.conf を環境変数から生成）
+        # rclone のセットアップ
         self.setup_rclone_conf()
 
         self.setup_routes()
@@ -69,7 +69,7 @@ class CivitAICrawler:
             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 からどうぞ。
+            モデル名とバックアップURLの紐づきは: https://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.log
             Status: {now} + currently running :D
             """
             return description
@@ -97,8 +97,12 @@ class CivitAICrawler:
 
     def encrypt_with_rclone(self, local_path: str):
         """
-        指定ファイル or ディレクトリを cryptLocal: にコピーし、
-        暗号化された結果を self.config.ENCRYPTED_DIR に生成する。
+        指定ファイル or ディレクトリを cryptLocal:にコピー。
+        このとき「トップレベルのフォルダ名」も暗号化して保持するため、
+        cryptLocal:{os.path.basename(local_path)} の形でコピーする。
+        
+        ※ rclone の crypt設定が filename_encryption = standard 等なら、
+          フォルダ名やファイル名も丸ごと暗号化される
         """
         if not os.path.exists(local_path):
             raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}")
@@ -107,10 +111,17 @@ class CivitAICrawler:
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
-        cmd = ["rclone", "copy", local_path, "cryptLocal:", "-v"]
+        # コピー先を cryptLocal:xxx に指定することで、xxx というフォルダ単位で
+        # 暗号化されたディレクトリを作成できるようにする。
+        top_level_name = os.path.basename(local_path.rstrip("/"))
+        if not top_level_name:
+            top_level_name = "unknown"
+
+        # 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)
-        logger.info(f"[OK] rclone copy => cryptLocal:")
+        logger.info(f"[OK] rclone copy => cryptLocal:{top_level_name}")
 
         if not os.path.isdir(self.config.ENCRYPTED_DIR):
             raise FileNotFoundError(
@@ -119,23 +130,25 @@ class CivitAICrawler:
 
     def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
         """
-        self.config.ENCRYPTED_DIR にある暗号化後のファイル/フォルダ構造をそっくりそのまま
-        Hugging Faceにアップロードする。フォルダ名・ファイル名はrcloneの設定によって暗号化されている。
+        self.config.ENCRYPTED_DIR にある暗号化後のファイル/フォルダ構造を
+        そのまま Hugging Face にアップロードする。
+        
+        rclone で filename_encryption=standard が有効な場合:
+         - 元フォルダ名/ファイル名は完全に暗号化され、HF上では判読不能な名称になる。
+        
+        base_path_in_repo が空("")の場合は、
+        `/home/user/app/encrypted` の構造が HF リポジトリの直下に展開される。
         """
         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 からの相対パスを算出
+                # self.config.ENCRYPTED_DIR からの相対パス（暗号化後のフォルダ名・ファイル名）
                 relative_path = os.path.relpath(encrypted_file_path, self.config.ENCRYPTED_DIR)
-                # 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)
 
                 attempt = 0
@@ -162,7 +175,7 @@ class CivitAICrawler:
                         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  # この場合はリトライ回数をカウントしない
+                            attempt -= 1
                         else:
                             if attempt < max_retries:
                                 logger.warning(
@@ -347,38 +360,37 @@ class CivitAICrawler:
         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
-            new_repo_id = re.sub(r'\d+$', str(number), repo_id)
+            return re.sub(r'\d+$', str(number), repo_id)
         else:
-            new_repo_id = f"{repo_id}1"
-        return new_repo_id
+            return f"{repo_id}1"
 
     # =============================================================================
-    # 既存のファイル／フォルダアップロードを「フォルダ名も暗号化」するよう変更
+    # フォルダアップロード
     # =============================================================================
     def upload_file(self, file_path: str):
         """
-        1) rcloneで file_path を暗号化
-        2) 暗号化された単一ファイル(またはフォルダ構造)をそのままHugging Faceへアップロード
-           （ファイル名もフォルダ名も暗号化される）
+        単一ファイルを暗号化 → 暗号化後のものをHugging Faceへアップロード。
+        rcloneでフォルダを作って暗号化するため、フォルダ構造としてアップされる。
         """
-        # rcloneで暗号化
-        self.encrypt_with_rclone(file_path)
-        # 暗号化フォルダ全体をアップロード (フォルダ名も暗号化される)
+        self.encrypt_with_rclone(file_path)  # 1) ローカルを -> cryptLocal:basename(…)
         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):
         """
-        1) rcloneでフォルダ全体を暗号化 (フォルダ名含むディレクトリ構造が暗号化)
-        2) 暗号化済みディレクトリ構造をHugging Faceにアップロード
+        フォルダを暗号化 → フォルダ名ごとアップロード。
+        rclone copy folder_path => cryptLocal:folder_pathのbasename
+        すると {ENCRYPTED_DIR}/{暗号化後のトップレベルフォルダ} が生成されるため、
+        それを再帰的にHugging Faceへアップロードする。
         """
         self.encrypt_with_rclone(folder_path)
         self.upload_encrypted_files(repo_id=self.repo_ids['current'], base_path_in_repo="")