Update main.py

main.py

@@ -46,7 +46,7 @@ class Config:
         "Content-Type": "application/json"
     }
 
-    # ===== rclone 用の追加設定 =====
+    # rclone 用の追加設定
     RCLONE_CONF_BASE64 = os.environ.get("RCLONE_CONF_BASE64", "")
     ENCRYPTED_DIR = "/home/user/app/encrypted"
 
@@ -61,9 +61,8 @@ class CivitAICrawler:
         self.repo_ids = self.config.REPO_IDS.copy()
         self.jst = self.config.JST
 
-        # rcloneのセットアップ
+        # rclone 設定の読み込み
         self.setup_rclone_conf()
-
         self.setup_routes()
 
     def setup_routes(self):
@@ -83,7 +82,7 @@ class CivitAICrawler:
             asyncio.create_task(self.crawl())
 
     # ============================================================================
-    # rclone設定 & 暗号化アップロード関連
+    # rclone 設定 & 暗号化アップロード処理
     # ============================================================================
     def setup_rclone_conf(self):
         if not self.config.RCLONE_CONF_BASE64:
@@ -97,10 +96,10 @@ class CivitAICrawler:
         logger.info(f"[INFO] rclone.conf created at: {conf_path}")
 
     def encrypt_with_rclone(self, local_path: str):
-        """単一ファイル or ディレクトリを cryptLocal: にコピーし、暗号化する"""
+        """フォルダ or ファイルを cryptLocal: にコピーし、フォルダ名・ファイル名を暗号化"""
         if not os.path.exists(local_path):
             raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}")
-        # 事前に暗号先ディレクトリをクリーンアップ
+        # 事前に暗号先ディレクトリを掃除
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
@@ -116,13 +115,14 @@ class CivitAICrawler:
             )
 
     def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
-        """self.config.ENCRYPTED_DIR 配下の暗号化済ファイルを再帰的にアップロード"""
+        """self.config.ENCRYPTED_DIR 以下の暗号化済ファイルを再帰的にアップロード"""
         max_retries = 5
         for root, dirs, files in os.walk(self.config.ENCRYPTED_DIR):
             for fn in files:
                 encrypted_file_path = os.path.join(root, fn)
                 if not os.path.isfile(encrypted_file_path):
                     continue
+
                 relative_path = os.path.relpath(encrypted_file_path, self.config.ENCRYPTED_DIR)
                 upload_path_in_repo = os.path.join(base_path_in_repo, relative_path)
 
@@ -139,7 +139,7 @@ class CivitAICrawler:
                     except Exception as e:
                         attempt += 1
                         error_message = str(e)
-                        # 429 Rate-limit (31 minutes)
+                        # 429 Rate-limit with "in XX minutes"
                         if "rate-limited" in error_message and "minutes" in error_message:
                             import re
                             match = re.search(r"in (\d+) minutes?", error_message)
@@ -149,13 +149,13 @@ class CivitAICrawler:
                                 time.sleep(minutes * 60)
                                 attempt -= 1
                                 continue
-                        # 1時間待機パターン
+                        # 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...")
                             time.sleep(3600)
                             attempt -= 1
                             continue
-                        # 100kファイル上限
+                        # ファイル上限
                         if "over the limit of 100000 files" in error_message:
                             logger.warning("Repository file limit exceeded. Creating a new repository...")
                             self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current'])
@@ -169,50 +169,52 @@ class CivitAICrawler:
                                 f"Failed to upload {encrypted_file_path}, retry {attempt}/{max_retries}..."
                             )
                         else:
-                            logger.error(
-                                f"Failed to upload after {max_retries} attempts: {encrypted_file_path}"
-                            )
+                            logger.error(f"Failed to upload after {max_retries} attempts: {encrypted_file_path}")
                             raise
 
-    @staticmethod
-    def increment_repo_name(repo_id: str) -> str:
-        match = re.search(r'(\d+)$', repo_id)
-        if match:
-            number = int(match.group(1)) + 1
-            return re.sub(r'\d+$', str(number), repo_id)
-        else:
-            return f"{repo_id}1"
+    def upload_folder_encrypted(self, folder_path: str, repo_id: Optional[str] = None, path_in_repo: str = ""):
+        """フォルダを丸ごと暗号化してアップロード (=フォルダ名も暗号化)"""
+        if not repo_id:
+            repo_id = self.repo_ids['current']
+
+        self.encrypt_with_rclone(folder_path)
+        self.upload_encrypted_files(repo_id, base_path_in_repo=path_in_repo)
+
+        # 暗号化フォルダを削除
+        if os.path.isdir(self.config.ENCRYPTED_DIR):
+            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
 
     # ============================================================================
-    # 単ファイル暗号化アップロード → ローカル削除
+    # 単一ファイルを暗号化アップロードしてローカル削除 (old_versions用)
     # ============================================================================
-    def upload_file_encrypted_one_by_one(
-        self,
-        local_path: str,
-        repo_id: Optional[str] = None,
-        path_in_repo: str = ""
-    ):
+    def upload_file_encrypted_one_by_one(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: str = ""):
         """
-        単一ファイル (or フォルダ) を暗号化してアップロードしたあと、ローカルファイルを削除する。
+        単一ファイルを暗号化アップロードし、アップロード後にローカルファイルを削除。
         """
         if not repo_id:
             repo_id = self.repo_ids['current']
 
-        self.encrypt_with_rclone(local_path)
-        self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=path_in_repo)
+        self.encrypt_with_rclone(file_path)
+        self.upload_encrypted_files(repo_id, base_path_in_repo=path_in_repo)
 
-        # 暗号化用のENCRYPTED_DIRを消す
+        # 暗号化ディレクトリを削除
         if os.path.isdir(self.config.ENCRYPTED_DIR):
             shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
+        # ローカルの実ファイル削除
+        if os.path.exists(file_path):
+            os.remove(file_path)
 
-        # 元のローカルファイルも削除
-        if os.path.isfile(local_path):
-            os.remove(local_path)
-        elif os.path.isdir(local_path):
-            shutil.rmtree(local_path, ignore_errors=True)
+    @staticmethod
+    def increment_repo_name(repo_id: str) -> str:
+        match = re.search(r'(\d+)$', repo_id)
+        if match:
+            number = int(match.group(1)) + 1
+            return re.sub(r'\d+$', str(number), repo_id)
+        else:
+            return f"{repo_id}1"
 
     # ============================================================================
-    # 生ファイルアップロード (ログなど)
+    # ログや model_list.log は生アップロード
     # ============================================================================
     def upload_file_raw(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None):
         if repo_id is None:
@@ -242,7 +244,7 @@ class CivitAICrawler:
                     repo_id = self.repo_ids['current']
                     continue
                 elif "you can retry this action in about 1 hour" in error_message:
-                    logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour...")
+                    logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour before retrying...")
                     time.sleep(3600)
                     attempt -= 1
                 else:
@@ -253,7 +255,7 @@ class CivitAICrawler:
                         raise
 
     # ============================================================================
-    # ダウンロード関連
+    # ダウンロード処理
     # ============================================================================
     @staticmethod
     def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
@@ -280,47 +282,67 @@ class CivitAICrawler:
                 file.write(chunk)
 
         logger.info(f"Download completed: {file_path}")
-        return file_path  # ★ ダウンロードしたファイルのパスを返すように
+        return file_path
 
     # ============================================================================
-    # （★修正）1ファイルずつDL→暗号化→アップロード→削除
+    # 古いバージョンのみ1ファイルずつアップロード
     # ============================================================================
-    def process_latest_files_one_by_one(self, version_data: dict, model_folder: str, encrypted_folder_name: str):
-        """
-        最新バージョンのファイルを1つずつダウンロード→暗号化アップロード→ローカル削除
-        path_in_repo は "{encrypted_folder_name}/" をベースに。
-        """
-        files = version_data.get("files", [])
-        for file_info in files:
+    def download_old_versions_one_by_one(self, version_list: list, folder: str):
+        """version_list[1:] を対象に、モデルファイルを 1ファイルDL→upload→削除 を繰り返す"""
+        if len(version_list) <= 1:
+            return
+
+        old_versions_folder = os.path.join(folder, "old_versions")
+        os.makedirs(old_versions_folder, exist_ok=True)
+
+        for version in version_list[1:]:
+            for file_info in version.get("files", []):
+                download_url = file_info["downloadUrl"]
+                file_name = file_info["name"]
+
+                local_path = self.download_file(download_url, old_versions_folder, file_name)
+                if not local_path or not os.path.exists(local_path):
+                    logger.error(f"Failed to download or file not found: {file_name}")
+                    continue
+
+                # 1つアップロードして削除
+                # path_in_repo を空文字にすればフォルダ名も暗号化される（トップレベル）
+                # もしサブフォルダにまとめたいなら "old_versions" とか指定する
+                self.upload_file_encrypted_one_by_one(local_path, path_in_repo="")
+
+        # old_versions フォルダ内は空になったはずなので削除
+        if os.path.exists(old_versions_folder):
+            shutil.rmtree(old_versions_folder, ignore_errors=True)
+
+    # ============================================================================
+    # 従来どおり「最新バージョンのファイル一式 + images」フォルダを一括DL→アップロード
+    # ============================================================================
+    def download_model(self, model_versions: list, folder: str):
+        """最新バージョンを一括ダウンロード (フォルダにまとめる)"""
+        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"]
+            local_path = self.download_file(download_url, folder, file_name)
+            if local_path and os.path.exists(local_path):
+                logger.info(f"Downloaded {file_name}")
+            else:
+                logger.warning(f"Could not download {file_name}")
 
-            # ダウンロード
-            local_path = self.download_file(download_url, model_folder, file_name)
-            if not local_path or not os.path.exists(local_path):
-                logger.warning(f"Skip because file not found locally: {local_path}")
-                continue
-
-            # 暗号化アップロード
-            # 例: "myModelName/filename"
-            in_repo_path = os.path.join(encrypted_folder_name, file_name)
-            self.upload_file_encrypted_one_by_one(local_path, repo_id=self.repo_ids['current'], path_in_repo=in_repo_path)
+    def download_images(self, model_versions: list, folder: str):
+        images_folder = os.path.join(folder, "images")
+        os.makedirs(images_folder, exist_ok=True)
 
-    def process_images_one_by_one(self, version_list: list, model_folder: str, encrypted_folder_name: str):
-        """
-        画像をすべて1つずつDL→暗号化アップロード→削除
-        path_in_repo は "{encrypted_folder_name}/images/"
-        """
         images = []
-        for version in version_list:
-            for img_info in version.get("images", []):
-                images.append(img_info["url"])
+        for version in model_versions:
+            for img in version.get("images", []):
+                images.append(img["url"])
 
         for image_url in images:
-            image_name = image_url.split("/")[-1] + ".png"
-            local_path = os.path.join(model_folder, image_name)
-
-            # ダウンロード
+            image_name = os.path.basename(image_url) + ".png"
+            local_path = os.path.join(images_folder, image_name)
             try:
                 resp = requests.get(image_url, stream=True)
                 resp.raise_for_status()
@@ -330,46 +352,7 @@ class CivitAICrawler:
                 logger.info(f"Downloaded image: {local_path}")
             except Exception as e:
                 logger.error(f"Error downloading image {image_url}: {e}")
-                continue
 
-            # アップロード
-            in_repo_path = os.path.join(encrypted_folder_name, "images", image_name)
-            self.upload_file_encrypted_one_by_one(local_path, self.repo_ids['current'], in_repo_path)
-
-    def process_old_versions_one_by_one(self, version_list: list, model_folder: str, encrypted_folder_name: str):
-        """
-        古いバージョン (index=1以降) のファイルを 1つずつダウンロード→暗号化アップロード→削除
-        path_in_repo は "{encrypted_folder_name}/old_versions/{versionID_orName}/filename"
-        """
-        if len(version_list) <= 1:
-            return
-
-        for old_version in version_list[1:]:
-            # どんな名前でフォルダを区別するか（バージョンIDやバージョン名など）
-            version_id_or_name = str(old_version.get("id", "old_ver"))
-            files = old_version.get("files", [])
-            for file_info in files:
-                download_url = file_info["downloadUrl"]
-                file_name = file_info["name"]
-
-                # ダウンロード
-                local_path = self.download_file(download_url, model_folder, file_name)
-                if not local_path or not os.path.exists(local_path):
-                    logger.warning(f"Skip because file not found locally: {local_path}")
-                    continue
-
-                # 暗号化アップロード
-                in_repo_path = os.path.join(
-                    encrypted_folder_name,
-                    "old_versions",
-                    version_id_or_name,
-                    file_name
-                )
-                self.upload_file_encrypted_one_by_one(local_path, self.repo_ids['current'], in_repo_path)
-
-    # ============================================================================
-    # HTML & model_info.json は軽量なので一括DL→アップロードでもOK
-    # ============================================================================
     def save_html_content(self, url: str, folder: str):
         try:
             response = requests.get(url)
@@ -377,20 +360,13 @@ class CivitAICrawler:
             html_path = os.path.join(folder, os.path.basename(folder) + ".html")
             with open(html_path, 'w', encoding='utf-8') as file:
                 file.write(response.text)
-            return html_path
         except Exception as e:
             logger.error(f"Error saving HTML content for URL {url}: {e}")
-            return None
 
-    def save_model_info(self, model_info: dict, folder: str):
-        json_path = os.path.join(folder, "model_info.json")
-        try:
-            with open(json_path, "w", encoding="utf-8") as file:
-                json.dump(model_info, file, indent=2)
-            return json_path
-        except Exception as e:
-            logger.error(f"Error saving model_info.json: {e}")
-            return None
+    @staticmethod
+    def save_model_info(model_info: dict, folder: str):
+        with open(os.path.join(folder, "model_info.json"), "w", encoding="utf-8") as file:
+            json.dump(model_info, file, indent=2)
 
     # ============================================================================
     # model_list.log
@@ -411,22 +387,23 @@ class CivitAICrawler:
             logger.error(f"Failed to read model list: {e}")
         return model_list
 
+    # ============================================================================
+    # model 情報取得
+    # ============================================================================
     def get_model_info(self, model_id: str) -> dict:
-        """
-        model_id（例: '1110807'）に対応するモデル情報を
-        CivitAIのAPIから取得し、jsonを返す
-        """
         try:
             url = self.config.URLS["modelId"] + str(model_id)
-            response = requests.get(url, headers=self.config.HEADERS)
-            response.raise_for_status()
-            return response.json()
+            resp = requests.get(url, headers=self.config.HEADERS)
+            resp.raise_for_status()
+            return resp.json()
         except requests.RequestException as e:
             logger.error(f"Failed to retrieve model info for ID {model_id}: {e}")
-            return {}  # or return None
-    
+            return {}
+
+    # ============================================================================
+    # メイン処理: 最新ファイル + images はフォルダごとアップロード。old_versions は1ファイルずつ。
+    # ============================================================================
     def process_model(self, model_url: str):
-        """ 指定されたモデルURLを処理 (1つずつファイルをDL→アップロード→削除) """
         try:
             model_id = model_url.rstrip("/").split("/")[-1]
             model_info = self.get_model_info(model_id)
@@ -434,103 +411,99 @@ class CivitAICrawler:
                 logger.error(f"No model_info returned for {model_id}")
                 return
 
-            latest_version = model_info.get("modelVersions", [])[0]
-            model_file = next(
-                (file for file in latest_version.get("files", []) if file.get("type") == "Model"),
-                None
-            )
+            model_versions = model_info.get("modelVersions", [])
+            if not model_versions:
+                logger.error(f"No modelVersions in model info {model_id}")
+                return
+
+            latest_version = model_versions[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"]
+                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"]
+                latest_filename = first_file['name']
                 folder = os.path.splitext(latest_filename)[0]
-                logger.warning(f"No 'Model' type file found for {model_id}. Using first file's name.")
+                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.log を読み込み、modelpage_name が既にあればスキップ)
+            # すでにアップ済みかどうか model_list.log でチェック (モデル名ベース)
             model_list = self.read_model_list()
             modelpage_name = model_info.get("name", f"Model_{model_id}")
             if modelpage_name in model_list.values():
                 logger.info(f"Model '{modelpage_name}' already in model_list. Skipping.")
-                # return  # 必要に応じてリターン
+                # 必要ならreturn
 
-            # HTMLやmodel_info は軽いので一括保存→まとめて単ファイル暗号化アップロード
-            html_path = self.save_html_content(self.config.URLS["modelPage"] + str(model_id), folder)
-            json_path = self.save_model_info(model_info, folder)
+            # 最新バージョン (まとめてダウンロード)
+            self.download_model(model_versions, folder)
 
-            # 暗号化アップロード（HTML, JSON など）
-            # HF 上では "{folder}/model_info.json" としておく例
-            if html_path and os.path.exists(html_path):
-                in_repo_path = os.path.join(folder, os.path.basename(html_path))
-                self.upload_file_encrypted_one_by_one(html_path, self.repo_ids['current'], in_repo_path)
+            # 画像 (imagesフォルダまるごとダウンロード)
+            self.download_images(model_versions, folder)
 
-            if json_path and os.path.exists(json_path):
-                in_repo_path = os.path.join(folder, "model_info.json")
-                self.upload_file_encrypted_one_by_one(json_path, self.repo_ids['current'], in_repo_path)
+            # HTML & model_info.json
+            self.save_html_content(self.config.URLS["modelPage"] + str(model_id), folder)
+            self.save_model_info(model_info, folder)
 
-            # 最新バージョンを1ファイルずつアップロード
-            self.process_latest_files_one_by_one(latest_version, folder, folder)
+            # 古いバージョンのみ「1つずつアップロード＆削除」
+            self.download_old_versions_one_by_one(model_versions, folder)
 
-            # 画像を1ファイルずつアップロード
-            self.process_images_one_by_one(model_info["modelVersions"], folder, folder)
+            # ↑で old_versions は空になった → あとはフォルダに残っているのは
+            #   最新バージョンファイル・imagesフォルダ・model_info.json・HTML など
 
-            # 古いバージョンを1ファイルずつアップロード
-            self.process_old_versions_one_by_one(model_info["modelVersions"], folder, folder)
+            # "folder" 自体を暗号化アップロード (= images フォルダごとアップロード)
+            #   path_in_repo を "" にすればフォルダ名も暗号化される
+            self.upload_folder_encrypted(folder, path_in_repo="")
 
-            # ここで folder はほぼ空だが、一応削除
+            # ローカルフォルダ削除
             if os.path.exists(folder):
                 shutil.rmtree(folder)
 
-            # 最後に model_list.log に追記 (「modelpage_name: HFのURL構造」)
-            # 今回はフォルダごと暗号化ではなくファイルごとなので、ひとまず "folder" をルート名に使っておく
-            # Hugging Face 上でのトップフォルダ URL:
-            #   https://huggingface.co/REPO_ID/tree/main/folder
-            model_hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{folder}"
+            # model_list.log に追記
+            # HF上では folder名 も暗号化されるが、ここでは元の "modelpage_name" と
+            # HFへのトップフォルダ参照URLを書く
+            model_hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main"
             with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
                 f.write(f"{modelpage_name}: {model_hf_url}\n")
 
         except Exception as e:
             logger.error(f"Unexpected error processing model ({model_url}): {e}")
 
-
+    # ============================================================================
+    # crawl
+    # ============================================================================
     async def crawl(self):
-        """新着モデルをチェックし、1��ずつ処理するループ"""
         while True:
             try:
                 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}")
 
                 local_file_path = hf_hub_download(repo_id=self.repo_ids["log"], filename=self.config.LOG_FILE)
                 shutil.copyfile(local_file_path, f"./{self.config.LOG_FILE}")
 
-                # civitai_backup.log を読み取り
+                # ログ読み込み
                 with open(self.config.LOG_FILE, "r", encoding="utf-8") as file:
                     lines = file.read().splitlines()
                     old_models = json.loads(lines[0]) if len(lines) > 0 else []
                     self.repo_ids["current"] = lines[1] if len(lines) > 1 else ""
 
-                # 新着モデルを確認
+                # 新着モデル確認
                 response = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                 response.raise_for_status()
                 latest_models = response.json().get("items", [])
-                latest_model_ids = [item.get("id") for item in latest_models if "id" in item]
+                latest_model_ids = [m["id"] for m in latest_models if "id" in m]
 
-                # 増分チェック
+                # 差集合
                 new_models = list(set(latest_model_ids) - set(old_models))
 
                 if new_models:
                     logger.info(f"New models found: {new_models}")
                     model_id = new_models[0]
 
-                    # 試行5回
                     for attempt in range(1, 6):
                         try:
                             self.process_model(f"{self.config.URLS['modelId']}{model_id}")
@@ -542,7 +515,7 @@ class CivitAICrawler:
                             else:
                                 await asyncio.sleep(2)
                 else:
-                    # 新モデルなし → backup.log を更新 & アップロード
+                    # 新モデルなし
                     with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                         f.write(json.dumps(latest_model_ids) + "\n")
                         f.write(f"{self.repo_ids['current']}\n")
@@ -555,14 +528,14 @@ class CivitAICrawler:
                     await asyncio.sleep(60)
                     continue
 
-                # 成功したモデルをold_modelsに追加 → backup.log更新
+                # 1件アップロードに成功したら old_models に追加
                 old_models.append(model_id)
                 with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                     f.write(json.dumps(old_models) + "\n")
                     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(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
                 self.upload_file_raw(self.config.LIST_FILE, self.repo_ids["model_list"], self.config.LIST_FILE)