Upload main.py

main.py

@@ -56,7 +56,6 @@ class CivitAICrawler:
 
         rclone_conf_base64 = os.environ.get("RCLONE_CONF_BASE64")
         if rclone_conf_base64:
-            # カレントディレクトリ配下に .rclone_config ディレクトリを作成
             config_dir = os.path.join(os.getcwd(), ".rclone_config")
             os.makedirs(config_dir, exist_ok=True)
 
@@ -64,7 +63,6 @@ class CivitAICrawler:
             with open(conf_path, "wb") as f:
                 f.write(base64.b64decode(rclone_conf_base64))
 
-            # rclone がここを参照するように設定
             os.environ["RCLONE_CONFIG"] = conf_path
             logger.info(f"[INFO] Created rclone.conf at {conf_path}")
         else:
@@ -115,7 +113,6 @@ class CivitAICrawler:
         filename = self.get_filename_from_cd(response.headers.get('content-disposition'), default_name)
         file_path = os.path.join(destination_folder, filename)
 
-        # ダウンロードとファイル保存処理
         with open(file_path, 'wb') as file:
             for chunk in response.iter_content(chunk_size=8192):
                 file.write(chunk)
@@ -136,10 +133,12 @@ class CivitAICrawler:
         """画像を images フォルダにまとめてダウンロードする."""
         images_folder = os.path.join(folder, "images")
         os.makedirs(images_folder, exist_ok=True)
+
         images = []
         for version in model_versions:
             for img in version.get("images", []):
                 images.append(img["url"])
+
         for image_url in images:
             image_name = os.path.basename(image_url)
             local_path = os.path.join(images_folder, image_name)
@@ -158,7 +157,7 @@ class CivitAICrawler:
         try:
             resp = requests.get(model_page_url)
             resp.raise_for_status()
-            html_path = os.path.join(folder, "page.html")
+            html_path = os.path.join(folder, f"{os.path.basename(folder)}.html")
             with open(html_path, 'w', encoding='utf-8') as f:
                 f.write(resp.text)
             logger.info(f"Saved HTML: {html_path}")
@@ -166,7 +165,7 @@ class CivitAICrawler:
             logger.error(f"Error saving HTML content from {model_page_url}: {e}")
 
     def save_model_info_json(self, model_info: dict, folder: str):
-        """モデル情報をJSONファイルとして保存."""
+        """モデル情報(json)の保存"""
         info_path = os.path.join(folder, "model_info.json")
         try:
             with open(info_path, 'w', encoding='utf-8') as f:
@@ -174,202 +173,50 @@ 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つのフォルダにダウンロードして
+        最後にまとめて暗号化アップロードする。
+        """
 
-        # 1) 最新バージョン (インデックス0) のファイルを folder にダウンロード
+        # 1) 最新バージョンは 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つずつダウンロード→暗号化アップロード→削除
+            local_path = self.download_file(download_url, folder, file_name)
+            # リトライ処理など省略
+
+        # 2) 古いバージョンを "folder/old_versions" にまとめる
         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))
-
-                        # ★ここでファイルのみを暗号化したいので encrypt_and_upload_single_file() を使用
-                        self.encrypt_and_upload_single_file(os.path.join(single_file_folder, file_name))
-                    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}")
-
-    def encrypt_and_upload_single_file(self, local_file: str) -> Optional[str]:
-        """
-        単一ファイルを暗号化→アップロード→ローカル削除する。
-          - rclone mkdir cryptLocal:enc_xxxx
-          - rclone copyto local_file => cryptLocal:enc_xxxx/filename
-          - 差分検知で作られた暗号フォルダを発見→upload_folder()→削除
-        """
-        import uuid
-        if not os.path.isfile(local_file):
-            logger.error(f"[encrypt_and_upload_single_file] Not a file: {local_file}")
-            return None
-
-        encrypted_base_dir = os.path.join(os.getcwd(), "encrypted")
-        os.makedirs(encrypted_base_dir, exist_ok=True)
-
-        # 古い暗号物を削除
-        before_dirs = set(os.listdir(encrypted_base_dir))
-        for item in before_dirs:
-            item_path = os.path.join(encrypted_base_dir, item)
-            try:
-                if os.path.isfile(item_path):
-                    os.remove(item_path)
-                else:
-                    shutil.rmtree(item_path)
-                logger.info(f"[CLEANUP] Removed old encrypted item: {item_path}")
-            except Exception as e:
-                logger.warning(f"[CLEANUP] Failed to remove {item_path}: {e}")
-
-        # 1) mkdir
-        subfolder_label = "enc_" + str(uuid.uuid4())[:8]
-        try:
-            subprocess.run(
-                ["rclone", "mkdir", f"cryptLocal:{subfolder_label}"],
-                check=True
-            )
-            logger.info(f"[OK] rclone mkdir cryptLocal:{subfolder_label}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"rclone mkdir failed: {e}")
-            return None
-
-        # 2) copyto (ファイル) => cryptLocal:enc_xxx/filename
-        filename_in_repo = os.path.basename(local_file)
-        try:
-            subprocess.run(
-                [
-                    "rclone", "copyto",
-                    local_file,
-                    f"cryptLocal:{subfolder_label}/{filename_in_repo}",
-                    "--create-empty-src-dirs"
-                ],
-                check=True
-            )
-            logger.info(f"[OK] rclone copyto {local_file} => cryptLocal:{subfolder_label}/{filename_in_repo}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"rclone copyto failed: {e}")
-            return None
-
-        # 3) 差分検知
-        after_dirs = set(os.listdir(encrypted_base_dir))
-        diff = after_dirs - before_dirs
-        if not diff:
-            logger.error("[ERROR] No new directory appeared in ./encrypted after rclone copyto.")
-            return None
-        if len(diff) > 1:
-            logger.warning(f"[WARN] Multiple new directories found in ./encrypted: {diff}")
-
-        enc_folder_name = diff.pop()  # 1つ抜き取る
-        enc_folder_path = os.path.join(encrypted_base_dir, enc_folder_name)
-        if not os.path.isdir(enc_folder_path):
-            logger.error(f"[ERROR] {enc_folder_path} is not a directory.")
-            return None
-
-        # 4) upload_folder
-        try:
-            self.upload_folder(enc_folder_path, path_in_repo=subfolder_label)
-            logger.info(f"[OK] Uploaded encrypted folder: {enc_folder_path}")
-        except Exception as e:
-            logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")
-
-        # 5) ローカル削除（暗号フォルダ + 元ファイル）
-        try:
-            if os.path.exists(local_file):
-                os.remove(local_file)
-            if os.path.isdir(enc_folder_path):
-                shutil.rmtree(enc_folder_path)
-            logger.info(f"[CLEANUP] Removed {local_file} and {enc_folder_path}")
-        except Exception as e:
-            logger.warning(f"[CLEANUP] Failed to remove local items: {e}")
-
-        return subfolder_label
+                    local_path = self.download_file(download_url, old_versions_folder, file_name)
 
     def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
         """
-        1. rclone mkdir cryptLocal:subfolder_label で空フォルダを作る
-        2. rclone copy local_folder => cryptLocal:subfolder_label --create-empty-src-dirs
-        3. そのフォルダを self.upload_folder() でアップ
-        4. ローカル平文フォルダ & 暗号フォルダを削除
-        5. subfolder_label を返す
+        1. rclone mkdir cryptLocal:subfolder_label
+        2. rclone copy local_folder => cryptLocal:subfolder_label
+        3. upload_folder() → ローカル削除
         """
-        if not os.path.exists(local_folder):
-            logger.error(f"encrypt_and_upload_folder: folder not found: {local_folder}")
+        if not os.path.isdir(local_folder):
+            logger.error(f"encrypt_and_upload_folder: {local_folder} is not a directory.")
             return None
 
         encrypted_base_dir = os.path.join(os.getcwd(), "encrypted")
         os.makedirs(encrypted_base_dir, exist_ok=True)
 
-        # 既存の暗号ファイル削除（不要ならコメントアウト）
+        # 既存の暗号物を削除
         for item in os.listdir(encrypted_base_dir):
             item_path = os.path.join(encrypted_base_dir, item)
             try:
-                if os.path.isfile(item_path) or os.path.islink(item_path):
+                if os.path.isfile(item_path):
                     os.remove(item_path)
                 else:
                     shutil.rmtree(item_path)
@@ -377,46 +224,48 @@ class CivitAICrawler:
             except Exception as e:
                 logger.warning(f"[CLEANUP] Failed to remove {item_path}: {e}")
 
-        # サブフォルダ名を生成（enc_ + UUID）
         subfolder_label = "enc_" + str(uuid.uuid4())[:8]
 
-        # ★ 追加: mkdir で先に空ディレクトリを作っておく
+        # mkdir
         try:
-            subprocess.run(
-                ["rclone", "mkdir", f"cryptLocal:{subfolder_label}"],
-                check=True
-            )
+            subprocess.run(["rclone", "mkdir", f"cryptLocal:{subfolder_label}"], check=True)
+            logger.info(f"[OK] rclone mkdir cryptLocal:{subfolder_label}")
         except subprocess.CalledProcessError as e:
             logger.error(f"rclone mkdir failed: {e}")
             return None
 
-        # ★ --create-empty-src-dirs オプションを付けて copy
+        # copy
         try:
             subprocess.run(
                 [
-                    "rclone", 
-                    "copy", 
-                    local_folder, 
+                    "rclone", "copy",
+                    local_folder,
                     f"cryptLocal:{subfolder_label}",
                     "--create-empty-src-dirs"
                 ],
                 check=True
             )
+            logger.info(f"[OK] rclone copy {local_folder} => cryptLocal:{subfolder_label}")
         except subprocess.CalledProcessError as e:
             logger.error(f"rclone copy failed: {e}")
             return None
 
-        # 暗号フォルダのパス（必ずディレクトリができている前提）
-        enc_folder_path = os.path.join(encrypted_base_dir, subfolder_label)
+        # 差分検知
+        after_list = os.listdir(encrypted_base_dir)
+        if not after_list:
+            logger.error("[ERROR] No new directory in ./encrypted after rclone copy.")
+            return None
+        enc_folder_name = after_list[0]  # 先頭を使う
+        enc_folder_path = os.path.join(encrypted_base_dir, enc_folder_name)
 
         if not os.path.isdir(enc_folder_path):
-            logger.error(f"[ERROR] {enc_folder_path} is not a directory, something is still off.")
+            logger.error(f"[ERROR] {enc_folder_path} is not a directory.")
             return None
 
-        # アップロード
+        # upload_folder
         try:
             self.upload_folder(enc_folder_path, path_in_repo=subfolder_label)
-            logger.info(f"Uploaded encrypted folder: {enc_folder_path}")
+            logger.info(f"[OK] Uploaded encrypted folder: {enc_folder_path}")
         except Exception as e:
             logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")
 
@@ -424,9 +273,9 @@ class CivitAICrawler:
         try:
             shutil.rmtree(local_folder)
             shutil.rmtree(enc_folder_path)
-            logger.info(f"Removed local folder: {local_folder} and {enc_folder_path}")
+            logger.info(f"[CLEANUP] Removed local folder: {local_folder} & {enc_folder_path}")
         except Exception as e:
-            logger.error(f"Failed to remove local folders: {e}")
+            logger.warning(f"[CLEANUP] Could not remove local folders: {e}")
 
         return subfolder_label
 
@@ -467,47 +316,7 @@ 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 リポジトリに一括アップロード
-        """
         if path_in_repo is None:
             path_in_repo = os.path.basename(folder_path)
 
@@ -544,7 +353,6 @@ class CivitAICrawler:
 
     @staticmethod
     def increment_repo_name(repo_id: str) -> str:
-        """リポジトリ名の末尾の数字をインクリメントする。"""
         match = re.search(r'(\d+)$', repo_id)
         if match:
             number = int(match.group(1)) + 1
@@ -554,7 +362,6 @@ class CivitAICrawler:
         return new_repo_id
 
     def read_model_list(self) -> dict:
-        """モデルリストを読み込む。"""
         model_list = {}
         try:
             with open(self.config.LIST_FILE, "r", encoding="utf-8") as f:
@@ -570,7 +377,6 @@ class CivitAICrawler:
         return model_list
 
     def get_repo_info(self, repo_id):
-        """リポジトリの情報を取得する。"""
         try:
             repo_info = self.api.repo_info(repo_id=repo_id, files_metadata=True)
             file_paths = [sibling.rfilename for sibling in repo_info.siblings]
@@ -580,41 +386,31 @@ class CivitAICrawler:
             return []
 
     def process_model(self, model_url: str):
-        """1つのモデルをダウンロードしてフォルダ丸ごと暗号化＆アップロードする."""
+        """1つのモデルをフォルダにまとめてダウンロード→暗号化→アップロード"""
         try:
-            # model_idを取得
             model_id = model_url.rstrip("/").split("/")[-1]
-
-            # モデル情報を取得
             model_info = self.get_model_info(model_id)
             if not model_info or "modelVersions" not in model_info:
                 logger.error(f"No valid model info for ID {model_id}. Skipping.")
                 return
 
-            # バージョン一覧
             versions = model_info["modelVersions"]
             if not versions:
                 logger.warning(f"No modelVersions found for ID {model_id}.")
                 return
 
-            # ベースとなるフォルダ名
-            folder_name = model_info.get("name", "UnknownModel")
+            folder_name = model_info.get("name", "UnnamedModel")
             folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)
             folder_name += "_" + str(uuid.uuid4())[:8]
             os.makedirs(folder_name, exist_ok=True)
 
-            # (A) 最新バージョンファイル + (B) 画像 + (C) HTML + model_info.json
-            #    → 全部 folder_name に置く
-            #    → 後でまとめて encrypt_and_upload_folder
+            # 最新 + 古いバージョンすべて folder_name 下へ
             self.download_and_process_versions(versions, folder_name)
             self.download_images(versions, folder_name)
             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]"
@@ -627,19 +423,18 @@ class CivitAICrawler:
             logger.error(f"Error in process_model ({model_url}): {e}")
 
     async def crawl(self):
-        """モデルを定期的にチェックし、更新を行う。"""
         while True:
             try:
                 login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
 
-                # model_list.logを最新化
+                # model_list.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
@@ -674,14 +469,12 @@ class CivitAICrawler:
                             else:
                                 await asyncio.sleep(2)
 
-                    # 新モデルを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 with new model ID: {model_id}")
 
-                    # ログファイル & model_list.logをアップロード
                     self.upload_file(
                         file_path=self.config.LOG_FILE,
                         repo_id=self.repo_ids["log"],
@@ -693,7 +486,6 @@ class CivitAICrawler:
                         path_in_repo=self.config.LIST_FILE
                     )
                 else:
-                    # 新着なし → ログを最新化してアップロードだけして待機
                     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")