Update main.py

main.py

@@ -233,9 +233,15 @@ class CivitAICrawler:
         except requests.RequestException as e:
             logger.error(f"Failed to retrieve model info for ID {model_id}: {e}")
 
-    def download_model(self, model_versions: list, folder: str, existing_old_version_files: list = []):
+    def download_model(self, model_versions: list, folder: str):
+        """
+        - 最新バージョンのファイルを「folder」へまとめてダウンロード
+        - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード)
+        （※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに）
+        """
         latest_version = model_versions[0]
         latest_files = latest_version["files"]
+
         for file_info in latest_files:
             download_url = file_info["downloadUrl"]
             file_name = file_info["name"]
@@ -267,45 +273,70 @@ class CivitAICrawler:
                 except Exception as e:
                     logger.error(f"Failed to create dummy file for {file_name}: {e}")
 
-        # 古いバージョンのダウンロード
-        if len(model_versions) > 1:
-            old_versions_folder = os.path.join(folder, "old_versions")
-            os.makedirs(old_versions_folder, exist_ok=True)
-            for version in model_versions[1:]:
-                for file_info in version["files"]:
-                    file_name = file_info["name"]
-                    if file_name in existing_old_version_files:
-                        logger.info(f"Skipping download of existing old version file: {file_name}")
-                        continue
-                    download_url = file_info["downloadUrl"]
-                    local_file_path = os.path.join(old_versions_folder, file_name)
-                    login_detected_count = 0
+    def download_old_version_file_and_upload(self, file_info, parent_folder: str, encrypted_top_name: str):
+        """
+        古いバージョンのファイルを1つダウンロード→暗号化アップロード→削除。
+        「old_versions」というフォルダ名をそのまま使う。
+        """
+        file_name = file_info["name"]
+        download_url = file_info["downloadUrl"]
 
-                    while login_detected_count < 5:
-                        try:
-                            self.download_file(download_url, old_versions_folder, file_name)
-                        except Exception as e:
-                            logger.error(f"Exception occurred while downloading {file_name}: {e}")
-                            login_detected_count += 1
-                            continue
+        # old_versions フォルダを作成
+        old_versions_folder = os.path.join(parent_folder, "old_versions")
+        os.makedirs(old_versions_folder, exist_ok=True)
 
-                        if "login" in os.listdir(old_versions_folder):
-                            login_detected_count += 1
-                            logger.warning(f"'login' file found while downloading {file_name}. Will try again. ({login_detected_count}/5)")
-                            os.remove(os.path.join(old_versions_folder, "login"))
-                        else:
-                            logger.info(f"Successfully downloaded {file_name}")
-                            break
+        local_path = os.path.join(old_versions_folder, file_name)
 
-                    if login_detected_count >= 5:
-                        dummy_file_name = f"{file_name}.download_failed"
-                        dummy_file_path = os.path.join(old_versions_folder, dummy_file_name)
-                        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_name}. URL: {download_url}")
-                        except Exception as e:
-                            logger.error(f"Failed to create dummy file for {file_name}: {e}")
+        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}")
+                login_detected_count += 1
+                continue
+
+            # 万が一 "login" という名前のファイルが出現したら削除して再試行
+            if "login" in os.listdir(old_versions_folder):
+                login_detected_count += 1
+                logger.warning(f"'login' file found while downloading {file_name}. Retry {login_detected_count}/5.")
+                os.remove(os.path.join(old_versions_folder, "login"))
+            else:
+                logger.info(f"Successfully downloaded old version file: {file_name}")
+                break
+
+        if login_detected_count >= 5:
+            # 5回失敗
+            dummy_file_name = f"{file_name}.download_failed"
+            dummy_file_path = os.path.join(old_versions_folder, dummy_file_name)
+            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: {dummy_file_name}")
+            except Exception as e:
+                logger.error(f"Failed to create dummy file for old version {file_name}: {e}")
+            return
+
+        # ダウンロード成功 → 1ファイルだけ暗号化＆アップロード → ローカル削除
+        try:
+            # 1) ファイル単位で暗号化するなら
+            self.encrypt_with_rclone(local_path)
+
+            # 2) 暗号化ファイルをアップロード
+            self.upload_encrypted_files(
+                repo_id=self.repo_ids["current"],
+                base_path_in_repo=encrypted_top_name
+            )
+            logger.info(f"Uploaded old version file: {file_name} into {encrypted_top_name}")
+
+        except Exception as e:
+            logger.error(f"Error uploading old version file {file_name}: {e}")
+
+        # 3) アップロード後、平文ファイルを削除
+        if os.path.exists(local_path):
+            os.remove(local_path)
+            logger.info(f"Removed local old version file: {local_path}")
 
     def download_images(self, model_versions: list, folder: str):
         images_folder = os.path.join(folder, "images")
@@ -475,12 +506,24 @@ class CivitAICrawler:
             return {}
 
     def process_model(self, model_url: str):
-        """指定されたモデルURLを処理する関数。"""
+        """
+        - 最新バージョンをフォルダに一括DL → フォルダごと暗号化アップロード
+        - 古いバージョンはファイル単位で即アップロード
+        """
         try:
             model_id = model_url.rstrip("/").split("/")[-1]
             model_info = self.get_model_info(model_id)
-    
-            latest_version = model_info.get("modelVersions", [])[0]
+
+            model_versions = model_info.get("modelVersions", [])
+            if not model_versions:
+                logger.warning(f"No versions found for model ID {model_id}")
+                return
+            
+            # ==================================================================
+            # 1) 「古いコード」と同じロジックでフォルダ名を決定
+            #    (最新バージョンのうち 'type' = 'Model' のファイルがあればそれ、なければ最初のファイル名)
+            # ==================================================================
+            latest_version = model_versions[0]
             model_file = next(
                 (file for file in latest_version["files"] if file.get('type') == 'Model'),
                 None
@@ -493,50 +536,46 @@ class CivitAICrawler:
                 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)
-    
-            # model_list を読み込み
-            model_list = self.read_model_list()
-    
-            # もし既に「同名（モデルページ名）がアップされている」かどうか確認したい場合の例:
-            #   ※ 今回は modelpage_name（= model_info["name"]） をキーにするか、
-            #     あるいは model_id (str) をキーにするか、運用に合わせて設定してください。
-            #   例として modelpage_name をキーとしてチェックする流れ:
-            modelpage_name = model_info.get("name", "Unnamed Model")
-    
-            if modelpage_name in model_list.values():
-                # 既に同モデルページ名がアップロード済み → ここでスキップや上書きなどの処理を決定
-                logger.info(f"Model '{modelpage_name}' is already listed in model_list. Skipping re-upload.")
-                # もし「強制再アップロード」したくないなら return で処理終了:
-                # return
-                # あるいは「強制アップするがバージョンだけ追加」などいろいろ処理が可能
-                # ここではあえて続行するが、必要に応じて書き換えてください。
-    
-            # ダウンロードや画像保存
-            existing_old_version_files = []
-            self.download_model(model_info["modelVersions"], folder, existing_old_version_files)
-            self.download_images(model_info["modelVersions"], folder)
+
+            # ==================================================================
+            # 2) 最新バージョンをまとめてダウンロード → フォルダごと暗号化アップロード
+            # ==================================================================
+            # 最新バージョンのファイルをまとめてダウンロード
+            self.download_model(model_versions, folder)
+
+            # 画像ダウンロード & HTML保存 & info.json保存
+            self.download_images(model_versions, folder)
             self.save_html_content(model_url, folder)
             self.save_model_info(model_info, folder)
-    
-            # ========== rclone で暗号化フォルダをアップロード ==========
+
+            # フォルダごと暗号化 → Hugging Face へアップ
             encrypted_top_name = self.upload_folder_encrypted(folder)
-    
-            # 今回アップロードしたモデルの URL
+            logger.info(f"[MAIN] Uploaded latest version folder => {encrypted_top_name}")
+
+            # ローカルフォルダ削除
+            shutil.rmtree(folder, ignore_errors=True)
+
+            # ==================================================================
+            # 3) 古いバージョンのファイルを1つずつDL → 同じ暗号化フォルダに追加
+            # ==================================================================
+            if len(model_versions) > 1:
+                for version in model_versions[1:]:
+                    for file_info in version["files"]:
+                        self.download_old_version_file_and_upload(file_info, encrypted_top_name)
+
+            # ==================================================================
+            # 4) 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}"
-    
-            # model_list.log に追記 → "modelpage_name: model_hf_url" 形式
+
             with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
                 f.write(f"{modelpage_name}: {model_hf_url}\n")
-    
-            # ローカルフォルダ削除
-            if os.path.exists(folder):
-                shutil.rmtree(folder)
-    
-        except Exception as e:
-            logger.error(f"Unexpected error processing model ({model_url}): {e}")
 
+        except Exception as e:
+            logger.error(f"Unexpected error in process_model ({model_url}): {e}")
 
     async def crawl(self):
         """モデルを定期的にチェックし、更新を行う。"""