Update main.py

main.py

@@ -233,17 +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):
         """
-        指定された model_versions (list) から
-        - 最新バージョン（先頭要素）のファイルを一括ダウンロード
-        - 古いバージョン (index=1以降) のファイルは1つずつダウンロード後に即アップロード・削除
+        - 最新バージョンのファイルを「folder」へまとめてダウンロード
+        - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード)
+        （※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに）
         """
-        # ----------------------
-        # 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"]
@@ -275,76 +273,67 @@ class CivitAICrawler:
                 except Exception as e:
                     logger.error(f"Failed to create dummy file for {file_name}: {e}")
 
-        # ----------------------
-        # 2) 古いバージョンのダウンロード → すぐアップロード＆削除
-        # ----------------------
-        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"]
-
-                    # 既にアップロード済み (existing_old_version_files) の場合はスキップ
-                    if file_name in existing_old_version_files:
-                        logger.info(f"Skipping download of existing old version file: {file_name}")
-                        continue
+    def download_old_version_file_and_upload(self, file_info, encrypted_top_name: str):
+        """
+        古いバージョンのファイルを1個ダウンロードしたら即アップロードし、ローカル削除。
+        - file_info: CivitAI API で得られる 'files' の要素(dict)
+        - encrypted_top_name: すでにアップ済みのフォルダのトップ (ランダム名)
+        """
+        file_name = file_info["name"]
+        download_url = file_info["downloadUrl"]
 
-                    download_url = file_info["downloadUrl"]
-                    local_file_path = os.path.join(old_versions_folder, file_name)
-                    login_detected_count = 0
+        # 一時フォルダ作成
+        temp_folder = "temp_old_file"
+        os.makedirs(temp_folder, exist_ok=True)
+        local_path = os.path.join(temp_folder, file_name)
 
-                    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
+        login_detected_count = 0
+        while login_detected_count < 5:
+            try:
+                self.download_file(download_url, temp_folder, file_name)
+            except Exception as e:
+                logger.error(f"Exception while downloading old file {file_name}: {e}")
+                login_detected_count += 1
+                continue
+
+            if "login" in os.listdir(temp_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(temp_folder, "login"))
+            else:
+                logger.info(f"Successfully downloaded old version file: {file_name}")
+                break
 
-                        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. "
-                                f"({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:
+            dummy_file_name = f"{file_name}.download_failed"
+            dummy_file_path = os.path.join(temp_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}")
+            # ダウンロード失敗したのでこのファイルはアップロードしない
+            shutil.rmtree(temp_folder, ignore_errors=True)
+            return
 
-                    # ダウンロードが5回失敗した場合はダミーを作って終了
-                    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}")
-                        # 失敗時はアップロードせずに次のファイルへ
-                        continue
+        # ダウンロード成功したので、暗号化してアップロード(既存のトップフォルダに追加)
+        try:
+            # encrypt_with_rclone(temp_folder) で temp_folder 全体を暗号化
+            # → upload_encrypted_files(... base_path_in_repo=encrypted_top_name) で
+            #    既存フォルダ(encrypted_top_name) 以下にファイルを追加
+            self.encrypt_with_rclone(temp_folder)
+            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}")
 
-                    # ダウンロードに成功した場合:
-                    # 1ファイルごとに暗号化してアップロードし、アップ後に削除
-                    try:
-                        # 例: path_in_repo はフォルダ構成したいなら "old_versions/..." とか
-                        # ここではフォルダ名=folder名/old_versions としておく
-                        self.upload_file_encrypted(
-                            file_path=local_file_path,
-                            path_in_repo=f"{folder}/old_versions/{file_name}"
-                        )
-                        # アップロードが問題なく完了したので削除
-                        os.remove(local_file_path)
-                        logger.info(f"Removed local old version file: {local_file_path}")
-                    except Exception as e:
-                        logger.error(f"Error uploading and removing old version file {local_file_path}: {e}")
-            # ### ここまで修正 ###
+        except Exception as e:
+            logger.error(f"Error uploading old version file {file_name}: {e}")
+
+        # 一時フォルダを削除
+        shutil.rmtree(temp_folder, ignore_errors=True)
 
     def download_images(self, model_versions: list, folder: str):
         images_folder = os.path.join(folder, "images")
@@ -514,68 +503,62 @@ 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_file = next(
-                (file for file in latest_version["files"] if file.get('type') == 'Model'),
-                None
-            )
-            if model_file:
-                latest_filename = model_file['name']
-                folder = os.path.splitext(latest_filename)[0]
-            else:
-                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.")
-    
+
+            # =====================================
+            # 1) フォルダ作成 & 最新バージョンDL
+            # =====================================
+            model_versions = model_info.get("modelVersions", [])
+            if not model_versions:
+                logger.warning(f"No versions found for model ID {model_id}")
+                return
+
+            # 例: 最新バージョン用フォルダ名を適当に
+            folder = f"model_{model_id}_latest"
             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)
+
+            # download_model() で最新バージョンだけダウンロード
+            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 で暗号化フォルダをアップロード ==========
+
+            # =====================================
+            # 2) 最新バージョンのフォルダを丸ごと暗号化 & アップロード
+            #    → 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):
         """モデルを定期的にチェックし、更新を行う。"""