Update main.py

main.py

@@ -233,15 +233,9 @@ 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):
-        """
-        - 最新バージョンのファイルを「folder」へまとめてダウンロード
-        - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード)
-        （※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに）
-        """
+    def download_model(self, model_versions: list, folder: str, existing_old_version_files: list = []):
         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"]
@@ -273,67 +267,45 @@ class CivitAICrawler:
                 except Exception as e:
                     logger.error(f"Failed to create dummy file for {file_name}: {e}")
 
-    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"]
-
-        # 一時フォルダ作成
-        temp_folder = "temp_old_file"
-        os.makedirs(temp_folder, exist_ok=True)
-        local_path = os.path.join(temp_folder, file_name)
-
-        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_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
+        # 古いバージョンのダウンロード
+        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
 
-        # ダウンロード成功したので、暗号化してアップロード(既存のトップフォルダに追加)
-        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}")
+                    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
 
-        except Exception as e:
-            logger.error(f"Error uploading old version file {file_name}: {e}")
+                        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
 
-        # 一時フォルダを削除
-        shutil.rmtree(temp_folder, ignore_errors=True)
+                    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}")
 
     def download_images(self, model_versions: list, folder: str):
         images_folder = os.path.join(folder, "images")
@@ -503,62 +475,68 @@ class CivitAICrawler:
             return {}
 
     def process_model(self, model_url: str):
-        """
-        - 最新バージョンをフォルダに一括DL → フォルダごと暗号化アップロード
-        - 古いバージョンはファイル単位で即アップロード
-        """
+        """指定されたモデルURLを処理する関数。"""
         try:
             model_id = model_url.rstrip("/").split("/")[-1]
             model_info = self.get_model_info(model_id)
-
-            # =====================================
-            # 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"
+    
+            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.")
+    
             os.makedirs(folder, exist_ok=True)
-
-            # download_model() で最新バージョンだけダウンロード
-            self.download_model(model_versions, folder)
-
-            # 画像ダウンロード & HTML保存 & info.json保存（従来通り）
-            self.download_images(model_versions, folder)
+    
+            # 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)
             self.save_html_content(model_url, folder)
             self.save_model_info(model_info, folder)
-
-            # =====================================
-            # 2) 最新バージョンのフォルダを丸ごと暗号化 & アップロード
-            #    → Hugging Face 側には"ランダムフォルダ名"が作られる
-            # =====================================
+    
+            # ========== rclone で暗号化フォルダをアップロード ==========
             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)
-
-            # =====================================
-            # 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}")
+    
+            # 今回アップロードしたモデルの URL
             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 in process_model ({model_url}): {e}")
+            logger.error(f"Unexpected error processing model ({model_url}): {e}")
+
 
     async def crawl(self):
         """モデルを定期的にチェックし、更新を行う。"""