Upload main.py

main.py

@@ -132,35 +132,6 @@ class CivitAICrawler:
             logger.error(f"Failed to retrieve model info for ID {model_id}: {e}")
             return {}
 
-    def download_model_files(self, model_versions: list, folder: str):
-        """最新のモデルバージョンと古いバージョンのファイルをまとめてダウンロード."""
-        for version in model_versions:
-            files_info = version.get("files", [])
-            for file_info in files_info:
-                download_url = file_info["downloadUrl"]
-                file_name = file_info["name"]
-                login_detected_count = 0
-
-                while login_detected_count < 5:
-                    local_path = self.download_file(download_url, folder, file_name)
-                    if local_path and "login" in os.listdir(folder):
-                        # 万が一、ダウンロード先に "login" という謎ファイルが出た場合の再試行処理
-                        login_detected_count += 1
-                        os.remove(os.path.join(folder, "login"))
-                        logger.warning(f"Detected 'login' file, retrying download: {file_name}  ({login_detected_count}/5)")
-                    else:
-                        break
-
-                if login_detected_count >= 5:
-                    # ダウンロード失敗を示すダミーファイルを作成
-                    dummy_file_path = os.path.join(folder, f"{file_name}.download_failed")
-                    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_path}")
-                    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 フォルダにまとめてダウンロードする."""
         images_folder = os.path.join(folder, "images")
@@ -203,20 +174,93 @@ class CivitAICrawler:
             logger.info(f"Saved model_info.json: {info_path}")
         except Exception as e:
             logger.error(f"Failed to save model info JSON: {e}")
+    # =========================================================================
+    # ここが重要:
+    #  - 最新バージョンはまとめて folder_name にダウンロード (一度に暗号化アップロード)
+    #  - 古いバージョンは1つずつダウンロード→暗号化→アップロード→削除 でストレージを節約
+    # =========================================================================
+    def download_and_process_versions(self, model_versions: list, folder: str):
+        """最新バージョンをまとめてダウンロード、old_versionsは1つずつアップして削除。"""
+
+        # 1) 最新バージョン (インデックス0) のファイルを folder にダウンロード
+        latest_version = model_versions[0]
+        logger.info(f"Processing latest version: {latest_version.get('name','(NoName)')}")
+
+        for file_info in latest_version.get("files", []):
+            download_url = file_info["downloadUrl"]
+            file_name = file_info["name"]
+            login_detected_count = 0
+
+            while login_detected_count < 5:
+                local_path = self.download_file(download_url, folder, file_name)
+                if local_path and "login" in os.listdir(folder):
+                    login_detected_count += 1
+                    os.remove(os.path.join(folder, "login"))
+                    logger.warning(f"'login' file found, retrying {file_name} ({login_detected_count}/5)")
+                else:
+                    break
+
+            if login_detected_count >= 5:
+                dummy_file_path = os.path.join(folder, f"{file_name}.download_failed")
+                try:
+                    with open(dummy_file_path, "w") as f:
+                        f.write("Download failed after 5 attempts.")
+                    logger.error(f"Failed to download {file_name}. Dummy file created: {dummy_file_path}")
+                except Exception as e:
+                    logger.error(f"Failed to create dummy file for {file_name}: {e}")
+
+        # 2) 古いバージョンがあれば 1つずつダウンロード→暗号化アップロード→削除
+        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:]:
+                logger.info(f"Processing older version: {version.get('name','(NoName)')}")
+                for file_info in version.get("files", []):
+                    file_name = file_info["name"]
+                    download_url = file_info["downloadUrl"]
+                    login_detected_count = 0
+
+                    while login_detected_count < 5:
+                        local_path = self.download_file(download_url, old_versions_folder, file_name)
+                        if local_path and "login" in os.listdir(old_versions_folder):
+                            login_detected_count += 1
+                            os.remove(os.path.join(old_versions_folder, "login"))
+                            logger.warning(f"'login' file found, retrying {file_name} ({login_detected_count}/5)")
+                        else:
+                            break
 
-
-    # =====================================================
-    # 必ずサブフォルダを作り、古いファイルやフォルダは削除してからコピー
-    # =====================================================
+                    if login_detected_count >= 5:
+                        dummy_file_path = os.path.join(old_versions_folder, f"{file_name}.download_failed")
+                        try:
+                            with open(dummy_file_path, "w") as f:
+                                f.write("Download failed after 5 attempts.")
+                            logger.error(f"Failed to download {file_name}. Dummy file: {dummy_file_path}")
+                        except Exception as e:
+                            logger.error(f"Failed to create dummy file for {file_name}: {e}")
+                        continue  # 次のファイルへ
+
+                    # ===== ダウンロード成功した古いバージョンファイルを暗号化アップロード =====
+                    # フォルダごと暗号化でもいいですが、大容量を避けるためファイル単位で暗号化する例
+                    # ここでは "encrypt_and_upload_folder" の代わりに「encrypt_and_upload_file」するなど
+                    #   あるいはフォルダごとでもOK
+                    single_file_folder = os.path.join(old_versions_folder, "temp_single")
+                    os.makedirs(single_file_folder, exist_ok=True)
+                    try:
+                        # 移動して「このファイルだけ」が入ったフォルダを作る
+                        single_file_path = shutil.move(local_path, os.path.join(single_file_folder, file_name))
+                        # 暗号化アップロード
+                        self.encrypt_and_upload_folder(single_file_folder)
+                    except Exception as e:
+                        logger.error(f"Failed to encrypt/upload old version file: {e}")
+                    finally:
+                        # single_file_folder削除（encrypt_and_upload_folderで消えるはず）
+                        if os.path.exists(single_file_folder):
+                            shutil.rmtree(single_file_folder)
+                            logger.info(f"Removed temp_single folder {single_file_folder}")
+
+    # ========== 既存の encrypt_and_upload_folder はそのまま活用 ==========
     def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
-        """
-        1. /home/user/app/encrypted 配下の古いファイルやフォルダを消去（オプション）
-        2. rclone copy local_folder => cryptLocal:subfolder_label
-           → /home/user/app/encrypted/subfolder_label/ が必ずフォルダとして作られる
-        3. そのフォルダを upload_folder()
-        4. ローカル平文フォルダ & 暗号フォルダを削除
-        5. 最後に subfolder_label を return
-        """
         if not os.path.exists(local_folder):
             logger.error(f"encrypt_and_upload_folder: folder not found: {local_folder}")
             return None
@@ -224,8 +268,7 @@ class CivitAICrawler:
         encrypted_base_dir = os.path.join(os.getcwd(), "encrypted")
         os.makedirs(encrypted_base_dir, exist_ok=True)
 
-        # --- 古い暗号ファイルやフォルダを消す例 ---
-        # これにより "/home/user/app/encrypted" 配下を毎回スッキリさせる
+        # 前のコードにある「古い暗号ファイル削除」はそのまま残す
         for item in os.listdir(encrypted_base_dir):
             item_path = os.path.join(encrypted_base_dir, item)
             try:
@@ -236,12 +279,8 @@ class CivitAICrawler:
                 logger.info(f"[CLEANUP] Removed old encrypted item: {item_path}")
             except Exception as e:
                 logger.warning(f"[CLEANUP] Failed to remove {item_path}: {e}")
-        # --------------------------------------
 
-        # 今回は "enc_{UUID}" のようなサブフォルダ名を決める
         subfolder_label = "enc_" + str(uuid.uuid4())[:8]
-
-        # rclone で確実に "enc_XXXX" フォルダが生成される
         try:
             subprocess.run(
                 ["rclone", "copy", local_folder, f"cryptLocal:{subfolder_label}"],
@@ -253,17 +292,17 @@ class CivitAICrawler:
 
         enc_folder_path = os.path.join(encrypted_base_dir, subfolder_label)
         if not os.path.isdir(enc_folder_path):
-            logger.error(f"[ERROR] {enc_folder_path} is not a directory. Rclone must have produced a file instead.")
+            logger.error(f"[ERROR] {enc_folder_path} is not a directory.")
             return None
 
-        # フォルダをアップロード
+        # アップロード
         try:
             self.upload_folder(enc_folder_path, path_in_repo=subfolder_label)
             logger.info(f"Uploaded encrypted folder: {enc_folder_path}")
         except Exception as e:
             logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")
 
-        # ローカル平文フォルダと暗号フォルダを削除
+        # ローカル削除
         try:
             shutil.rmtree(local_folder)
             shutil.rmtree(enc_folder_path)
@@ -310,6 +349,43 @@ class CivitAICrawler:
                         logger.error(f"Failed after {max_retries} attempts: {e}")
                         raise
 
+    def upload_file(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']
+        if path_in_repo is None:
+            path_in_repo = os.path.basename(file_path)
+        
+        max_retries = 5
+        attempt = 0
+        while attempt < max_retries:
+            try:
+                self.api.upload_file(
+                    path_or_fileobj=file_path,
+                    repo_id=repo_id,
+                    path_in_repo=path_in_repo
+                )
+                logger.info(f"Uploaded file: {file_path} to {repo_id} at {path_in_repo}")
+                return
+            except Exception as e:
+                attempt += 1
+                error_message = str(e)
+                if "over the limit of 100000 files" in error_message:
+                    logger.warning("File limit exceeded, creating a new repo.")
+                    self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current'])
+                    self.api.create_repo(repo_id=self.repo_ids['current'], private=True)
+                    attempt = 0
+                    continue
+                elif "you can retry this action in about 1 hour" in error_message:
+                    logger.warning("Rate limit hit. Waiting 1 hour...")
+                    time.sleep(3600)
+                    attempt -= 1
+                else:
+                    if attempt < max_retries:
+                        logger.warning(f"Failed to upload {file_path}, retry {attempt}/{max_retries}")
+                    else:
+                        logger.error(f"Failed after {max_retries} attempts: {e}")
+                        raise
+
     def upload_folder(self, folder_path: str, path_in_repo: Optional[str] = None):
         """
         フォルダを Hugging Face リポジトリに一括アップロード
@@ -403,36 +479,30 @@ class CivitAICrawler:
                 logger.warning(f"No modelVersions found for ID {model_id}.")
                 return
 
-            # フォルダ名として適当な名前をつける
-            # たとえばモデル名をベースにフォルダを作る(被り防止にUUIDを付与)
+            # ベースとなるフォルダ名
             folder_name = model_info.get("name", "UnknownModel")
-            folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)  # OSで使えない文字除去
+            folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)
             folder_name += "_" + str(uuid.uuid4())[:8]
             os.makedirs(folder_name, exist_ok=True)
 
-            # モデルファイルをダウンロード
-            self.download_model_files(versions, folder_name)
-
-            # 画像を images/ にダウンロード
+            # (A) 最新バージョンファイル + (B) 画像 + (C) HTML + model_info.json
+            #    → 全部 folder_name に置く
+            #    → 後でまとめて encrypt_and_upload_folder
+            self.download_and_process_versions(versions, folder_name)
             self.download_images(versions, folder_name)
-
-            # HTMLを取得
-            model_page_url = f"{self.config.URLS['modelPage']}{model_id}"
-            self.save_html_content(model_page_url, folder_name)
-
-            # model_info.json保存
+            self.save_html_content(f"{self.config.URLS['modelPage']}{model_id}", folder_name)
             self.save_model_info_json(model_info, folder_name)
 
-            # === フォルダごと暗号化 & アップロード
+            # ↑ この時点で「最新バージョン」の大ファイルが folder_name に残っている
+            #   しかし old_versions は既に1つずつ暗号化→削除済み
+
+            # まとめて暗号化 & アップロード
             enc_subfolder = self.encrypt_and_upload_folder(folder_name)
             if enc_subfolder is None:
                 enc_subfolder = "[ENCRYPT_FAILED]"
 
             hf_enc_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{enc_subfolder}"
-
-            # model_list.logに追記 (実際の暗号フォルダ名をそのまま書き込む)
             with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
-                # 例: "Pudu chileno (ID:12345): https://huggingface.co/xxx/tree/main/<暗号フォルダ>"
                 f.write(f"{model_info.get('name', 'UnnamedModel')} (ID:{model_id}): {hf_enc_url}\n")
 
         except Exception as e: