Upload main.py

main.py

@@ -104,9 +104,8 @@ class CivitAICrawler:
 
     def encrypt_with_rclone(self, local_path: str):
         """
-        local_path がファイルかフォルダかで処理を分け、
-        - フォルダ → 従来通り rclone copy
-        - ファイル → rclone copyto を使って暗号化
+        local_path がフォルダかファイルかで処理を分け、暗号化（rclone copy or copyto）を行う。
+        暗号化先は self.config.ENCRYPTED_DIR に生成される。
         """
         if not os.path.exists(local_path):
             raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}")
@@ -147,8 +146,16 @@ class CivitAICrawler:
                 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 = ""):
+        """
+        暗号化後の self.config.ENCRYPTED_DIR を走査し、
+        すべてのファイルを1つずつアップロードする。
+        
+        （「最新バージョンのフォルダをまとめてアップロードするとき」は、1回だけこのメソッドを呼び出して
+         フォルダ分のファイルを一括アップロード。
+         old_versions でも 1ファイル暗号化後にこのメソッドを呼び出しているため、
+         そこだけは引き続き個別ファイルのアップロードとなる。）
+        """
         max_retries = 5
         for root, dirs, files in os.walk(self.config.ENCRYPTED_DIR):
             for fn in files:
@@ -174,30 +181,21 @@ class CivitAICrawler:
                         attempt += 1
                         error_message = str(e)
 
-                        # ================================
-                        # 429によるrate-limit検出追加
-                        # ================================
-                        # "You have been rate-limited; you can retry this action in 31 minutes."
-                        # のようなメッセージから時間を抽出し、その時間+1分だけ待機後、再試行
+                        # 429制限等のリトライロジック
                         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 カウントで再試行
+                                attempt -= 1
                                 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  # 再試行回数を増やさずにループを続ける
+                            attempt -= 1
                             continue
 
                         if "over the limit of 100000 files" in error_message:
@@ -255,8 +253,7 @@ class CivitAICrawler:
     def download_model(self, model_versions: list, folder: str):
         """
         - 最新バージョンのファイルを「folder」へまとめてダウンロード
-        - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード)
-        （※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに）
+        - ここでは古いバージョンはダウンロードしない（後ほど old_versions にて個別にダウンロード）
         """
         latest_version = model_versions[0]
         latest_files = latest_version["files"]
@@ -295,12 +292,12 @@ class CivitAICrawler:
     def download_old_version_file_and_upload(self, file_info, parent_folder: str, encrypted_top_name: str):
         """
         古いバージョンのファイルを1つダウンロード→暗号化アップロード→削除。
-        「old_versions」というフォルダ名をそのまま使う。
+        old_versions フォルダの中に都度ファイルを落としてアップロード→削除する。
         """
         file_name = file_info["name"]
         download_url = file_info["downloadUrl"]
 
-        # old_versions フォルダを作成
+        # old_versions フォルダ作成
         old_versions_folder = os.path.join(parent_folder, "old_versions")
         os.makedirs(old_versions_folder, exist_ok=True)
 
@@ -309,7 +306,6 @@ class CivitAICrawler:
         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}")
@@ -326,7 +322,7 @@ class CivitAICrawler:
                 break
 
         if login_detected_count >= 5:
-            # 5回失敗
+            # 5回失敗した場合はダミーファイル化
             dummy_file_name = f"{file_name}.download_failed"
             dummy_file_path = os.path.join(old_versions_folder, dummy_file_name)
             try:
@@ -337,12 +333,12 @@ class CivitAICrawler:
                 logger.error(f"Failed to create dummy file for old version {file_name}: {e}")
             return
 
-        # ダウンロード成功 → 1ファイルだけ暗号化＆アップロード → ローカル削除
+        # ここから暗号化→アップロード→削除 （1ファイルずつ処理を繰り返す）
         try:
-            # 1) ファイル単位で暗号化するなら
+            # 1) 単ファイルを暗号化
             self.encrypt_with_rclone(local_path)
 
-            # 2) 暗号化ファイルをアップロード
+            # 2) アップロード
             self.upload_encrypted_files(
                 repo_id=self.repo_ids["current"],
                 base_path_in_repo=f"{encrypted_top_name}/old_versions"
@@ -352,7 +348,7 @@ class CivitAICrawler:
         except Exception as e:
             logger.error(f"Error uploading old version file {file_name}: {e}")
 
-        # 3) アップロード後、平文ファイルを削除
+        # 3) ローカルの平文ファイルを削除
         if os.path.exists(local_path):
             os.remove(local_path)
             logger.info(f"Removed local old version file: {local_path}")
@@ -402,7 +398,7 @@ class CivitAICrawler:
             return f"{repo_id}1"
 
     # =============================================================================
-    # 暗号化しないアップロード（ログや model_list.log 用）
+    # ログ類のアップロード（生ファイル）
     # =============================================================================
     def upload_file_raw(
         self,
@@ -448,26 +444,7 @@ class CivitAICrawler:
                         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,
@@ -475,12 +452,18 @@ class CivitAICrawler:
         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 ""
 
+        # 1) 全体を一度に暗号化
         self.encrypt_with_rclone(folder_path)
 
+        # 2) 暗号化後フォルダ名を取得し、一括アップロード
         top_levels = [
             d for d in os.listdir(self.config.ENCRYPTED_DIR)
             if os.path.isdir(os.path.join(self.config.ENCRYPTED_DIR, d))
@@ -491,56 +474,35 @@ class CivitAICrawler:
             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)
 
+        # 3) 暗号化用ディレクトリをクリーンアップ
         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」で扱うよう変更
+    # メイン処理：最新バージョンは「フォルダ一括」アップロード、old_versions は1ファイルずつ
     # =============================================================================
-    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 → 即暗号化アップロード
+        1) 最新バージョン → フォルダにまとめてダウンロードし、1回だけ暗号化&アップロード
+        2) old_versions → 1ファイルずつダウンロード→暗号化アップロード→削除
         """
         try:
             model_id = model_url.rstrip("/").split("/")[-1]
             model_info = self.get_model_info(model_id)
+            if not model_info:
+                logger.warning(f"Cannot get model info: {model_id}")
+                return
 
             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'),
@@ -550,48 +512,37 @@ class CivitAICrawler:
                 latest_filename = model_file['name']
                 folder = os.path.splitext(latest_filename)[0]
             else:
+                # type='Model' がなければ最初のファイルを基準に
                 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_html_content(self.config.URLS["modelPage"] + str(model_id), 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 → 即アップロード
-            # ==================================================================
+            # === 古いバージョンは1つずつアップロード ===
             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)
+                        # old_versions フォルダに落としてアップロード
                         self.download_old_version_file_and_upload(file_info, folder, encrypted_top_name)
-                        # ↑ parent_folder=folder に old_versionsフォルダを作り
-                        #    そこへダウンロードしてから暗号化アップロード
 
-            # ==================================================================
-            # 5) model_list.log への登録やその他の処理
-            # ==================================================================
+            # === 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}"
 
@@ -601,13 +552,17 @@ class CivitAICrawler:
         except Exception as e:
             logger.error(f"Unexpected error in process_model ({model_url}): {e}")
 
+    # =============================================================================
+    # 定期巡回して新しいモデルがあれば処理する
+    # =============================================================================
     async def crawl(self):
         """モデルを定期的にチェックし、更新を行う。"""
         while True:
             try:
+                # HuggingFaceログイン
                 login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
 
-                # model_list.log & civitai_backup.log を取得
+                # 既存の 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}")
 
@@ -670,7 +625,7 @@ class CivitAICrawler:
                     f.write(f"{self.repo_ids['current']}\n")
                 logger.info(f"Updated log file with new model ID: {model_id}")
 
-                # ログとmodel_list.logをアップロード
+                # ログとmodel_list.logをアップロード (暗号化しない)
                 self.upload_file_raw(
                     file_path=self.config.LOG_FILE,
                     repo_id=self.repo_ids["log"],