Upload main.py

main.py

@@ -7,7 +7,6 @@ import re
 import shutil
 import subprocess
 import time
-import uuid
 from typing import Optional
 
 import requests
@@ -25,9 +24,9 @@ class Config:
     LOG_FILE = "civitai_backup.log"
     LIST_FILE = "model_list.log"
     REPO_IDS = {
-        "log": "ttttdiva/CivitAI_log_test",       # ←ログ用リポジトリ
-        "model_list": "ttttdiva/CivitAI_model_info_test",  # ←モデル一覧用リポジトリ
-        "current": ""  # ←実際のアップ先
+        "log": "ttttdiva/CivitAI_log_test",       
+        "model_list": "ttttdiva/CivitAI_model_info_test",
+        "current": ""  
     }
     URLS = {
         "latest": "https://civitai.com/api/v1/models?sort=Newest",
@@ -81,10 +80,7 @@ class CivitAICrawler:
         async def startup_event():
             asyncio.create_task(self.crawl())
 
-    # ここで「download_file」や「encrypt_and_upload_folder」等のサポート関数を定義
-
     def download_file(self, url: str, dest_folder: str, filename: str) -> Optional[str]:
-        """実際にファイルをダウンロードし、dest_folder/filename に保存。"""
         os.makedirs(dest_folder, exist_ok=True)
         try:
             r = requests.get(url, headers=self.config.HEADERS, stream=True)
@@ -101,7 +97,6 @@ class CivitAICrawler:
         return file_path
 
     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:
@@ -118,7 +113,6 @@ class CivitAICrawler:
             logger.error(f"[ERR] upload_file: {e}")
 
     def upload_folder(self, folder_path: str, path_in_repo: Optional[str] = None):
-        """フォルダ丸ごとアップロード"""
         if path_in_repo is None:
             path_in_repo = os.path.basename(folder_path)
 
@@ -133,98 +127,61 @@ class CivitAICrawler:
             logger.error(f"[ERR] upload_folder: {e}")
 
     def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
-        """
-        1) subfolder_label = モデルフォルダ名 (local_folderのベース名)
-        2) rclone mkdir cryptLocal:subfolder_label
-        3) rclone copy local_folder => cryptLocal:subfolder_label
-        4) rclone copy cryptLocal:subfolder_label => encrypted/subfolder_label (差分検知やめる)
-        5) upload_folder(encrypted/subfolder_label, path_in_repo=subfolder_label)
-        6) cleanup
-        """
+        """local_folder -> cryptLocal: => encrypted/??? => upload_folder => cleanup"""
         if not os.path.isdir(local_folder):
             logger.error(f"[ERR] {local_folder} is not a directory.")
             return None
 
-        subfolder_label = os.path.basename(local_folder)  # ★ モデル名 or ファイル名を取得
-        if not subfolder_label:
-            logger.error(f"[ERR] local_folder has empty basename: {local_folder}")
-            return None
+        encrypted_dir = os.path.join(os.getcwd(), "encrypted")
+        os.makedirs(encrypted_dir, exist_ok=True)
 
-        # === 1) rclone mkdir
-        mkdir_cmd = ["rclone", "mkdir", f"cryptLocal:{subfolder_label}"]
-        logger.info(f"[CMD] {' '.join(mkdir_cmd)}")
+        before = set(os.listdir(encrypted_dir))
+        cmd = ["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"]
+        logger.info(f"[CMD] {' '.join(cmd)}")
         try:
-            subprocess.run(mkdir_cmd, check=True)
-            logger.info(f"[OK] rclone mkdir cryptLocal:{subfolder_label}")
+            subprocess.run(cmd, check=True)
+            logger.info("[OK] rclone copy => cryptLocal:")
         except subprocess.CalledProcessError as e:
-            logger.error(f"[ERR] rclone mkdir => {e}")
+            logger.error(f"[ERR] rclone copy failed: {e}")
             return None
 
-        # === 2) rclone copy local_folder => cryptLocal:subfolder_label
-        copy_cmd = ["rclone", "copy", local_folder, f"cryptLocal:{subfolder_label}", "--create-empty-src-dirs"]
-        logger.info(f"[CMD] {' '.join(copy_cmd)}")
-        try:
-            subprocess.run(copy_cmd, check=True)
-            logger.info(f"[OK] rclone copy => cryptLocal:{subfolder_label}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"[ERR] rclone copy => {e}")
+        after = set(os.listdir(encrypted_dir))
+        diff = after - before
+        if not diff:
+            logger.error("[ERR] no new directory in ./encrypted after copy")
             return None
-
-        # === 3) ローカル "encrypted/subfolder_label" に復号ファイル(?) を落とすかどうか
-        #     → HFへのフォルダアップには「ローカルに暗号ファイル群」が必要
-        #     → しかし "directory_name_encryption=true" ならリモート名だけで復号されず、ランダム名かも
-        # 今回は「物理的に local へ落として upload_folder()」する想定にする
-        encrypted_base = os.path.join(os.getcwd(), "encrypted")
-        os.makedirs(encrypted_base, exist_ok=True)
-
-        # remove old if exist
-        local_enc_path = os.path.join(encrypted_base, subfolder_label)
-        if os.path.exists(local_enc_path):
-            shutil.rmtree(local_enc_path)
-
-        # rclone copy cryptLocal:subfolder_label => encrypted/subfolder_label
-        pull_cmd = ["rclone", "copy", f"cryptLocal:{subfolder_label}", local_enc_path, "--create-empty-src-dirs"]
-        logger.info(f"[CMD] {' '.join(pull_cmd)}")
-        try:
-            subprocess.run(pull_cmd, check=True)
-            logger.info(f"[OK] rclone copy => {local_enc_path}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"[ERR] rclone copy back => {e}")
+        if len(diff) > 1:
+            logger.warning(f"[WARN] multiple new dirs => {diff}")
+        enc_name = diff.pop()
+        enc_path = os.path.join(encrypted_dir, enc_name)
+        if not os.path.isdir(enc_path):
+            logger.error(f"[ERR] {enc_path} is not a directory.")
             return None
 
-        # === 4) upload_folder(local_enc_path, path_in_repo=subfolder_label)
+        # HF upload folder
         try:
-            self.upload_folder(local_enc_path, path_in_repo=subfolder_label)
+            self.upload_folder(enc_path, path_in_repo=enc_name)
         except Exception as e:
             logger.error(f"[ERR] encrypt_and_upload_folder => upload_folder: {e}")
 
-        # === 5) cleanup
-        # remove local_folder & local_enc_path
+        # cleanup
         shutil.rmtree(local_folder, ignore_errors=True)
-        shutil.rmtree(local_enc_path, ignore_errors=True)
-        logger.info(f"[CLEANUP] removed {local_folder} & {local_enc_path}")
+        shutil.rmtree(enc_path, ignore_errors=True)
+        logger.info(f"[CLEANUP] removed {local_folder} & {enc_path}")
 
-        return subfolder_label
+        # ★ 成功したら enc_name を返す
+        return enc_name
 
-    # ここで「download_and_process_versions」(＝古いバージョン含めまとめてダウンロード) 定義
-    #  => 今回は名前を合わせ「download_and_process_versions」に合わせる
-    
     def download_and_process_versions(self, model_versions: list, folder: str):
-        """
-        例: 最新バージョン => folder
-            古いバージョン => folder/old_versions
-        """
         if not model_versions:
             return
 
-        # 最新
         latest = model_versions[0]
         for f_info in latest.get("files", []):
             url = f_info["downloadUrl"]
             fname = f_info["name"]
             self.download_file(url, folder, fname)
 
-        # 古い
         if len(model_versions) > 1:
             ov_folder = os.path.join(folder, "old_versions")
             os.makedirs(ov_folder, exist_ok=True)
@@ -235,7 +192,6 @@ class CivitAICrawler:
                     self.download_file(url, ov_folder, fname)
 
     def get_model_info(self, model_id: str) -> dict:
-        """モデルIDからCivitAIの情報を取得"""
         try:
             url = f"{self.config.URLS['modelId']}{model_id}"
             resp = requests.get(url, headers=self.config.HEADERS)
@@ -245,6 +201,20 @@ class CivitAICrawler:
             logger.error(f"[ERR] get_model_info({model_id}): {e}")
             return {}
 
+    def download_images(self, model_versions: list, folder: str):
+        """
+        各model_versionsから画像URLを集めて、
+        folder/images 下にダウンロードするメソッド。
+        """
+        images_folder = os.path.join(folder, "images")
+        os.makedirs(images_folder, exist_ok=True)
+
+        for version in model_versions:
+            for img_info in version.get("images", []):
+                img_url = img_info["url"]
+                filename = os.path.basename(img_url)
+                self.download_file(img_url, images_folder, filename)
+
     def process_model(self, model_id: str):
         info = self.get_model_info(model_id)
         if not info or "modelVersions" not in info:
@@ -252,57 +222,57 @@ class CivitAICrawler:
             return
         
         versions = info["modelVersions"]
-
-        base_dir = "local_models" 
+        base_dir = "local_models"
         os.makedirs(base_dir, exist_ok=True)
 
-        # フォルダ名はモデル名にする場合: 
-        # name = info.get("name","UnknownModel")
-        # folder_name = re.sub(r'[\\/*?:"<>|]', '_', name)
-        # ただし今はIDが使われているので:
-        folder_name = str(model_id)
-
-        folder_path = os.path.join(base_dir, folder_name)
+        # モデル名
+        model_name = info.get("name", f"ID_{model_id}")
+        safe_name = re.sub(r'[\\/*?:"<>|]', '_', model_name)  # OSで使えない文字を _
+        folder_path = os.path.join(base_dir, safe_name)
         if os.path.exists(folder_path):
             shutil.rmtree(folder_path)
         os.makedirs(folder_path, exist_ok=True)
-
         logger.info(f"[OK] Created local folder => {folder_path}")
 
         # ダウン���ード
-        for v in versions:
-            for f_info in v.get("files", []):
-                dl_url = f_info["downloadUrl"]
-                fname = f_info["name"]
-                self.download_file(dl_url, folder_path, fname)
-        
-        # 画像
+        self.download_and_process_versions(versions, folder_path)
         self.download_images(versions, folder_path)
 
-        # OPTIONAL: save model_info.json or HTML
-        # self.save_model_info(info, folder_path)
-
         # === 暗号化＆アップロード ===
-        logger.info(f"[DEBUG] Calling encrypt_and_upload_folder => {folder_path}")
-        enc_folder = self.encrypt_and_upload_folder(folder_path)
-        if enc_folder is None:
-            logger.warning("[WARN] Encryption/Upload returned None.")
+        logger.info(f"[DEBUG] encrypt_and_upload_folder => {folder_path}")
+        enc_subfolder = self.encrypt_and_upload_folder(folder_path)
+        if enc_subfolder is None:
+            # 失敗
+            enc_subfolder = "[ENCRYPT_FAILED]"
         else:
-            logger.info(f"[OK] Encrypted & uploaded folder => {enc_folder}")
+            logger.info(f"[OK] Encrypted & uploaded => {enc_subfolder}")
+
+        # ★ model_list.log に追記
+        #   enc_subfolder が "[ENCRYPT_FAILED]" でなければ成功したフォルダ名
+        hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{enc_subfolder}"
+        model_list_line = f"{model_name} (ID:{model_id}): {hf_url}\n"
+        try:
+            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
+                f.write(model_list_line)
+            logger.info(f"[OK] Wrote to model_list.log => {model_list_line.strip()}")
+        except Exception as e:
+            logger.error(f"[ERR] writing model_list.log => {e}")
+
+        # ★ model_list.log をアップロード
+        self.upload_file(self.config.LIST_FILE, self.repo_ids["model_list"], self.config.LIST_FILE)
 
     async def crawl(self):
         while True:
             try:
                 login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
-                # model_list, log_file のダウンロード
-                mlist_path = hf_hub_download(repo_id=self.repo_ids["model_list"], filename=self.config.LIST_FILE)
-                shutil.copyfile(mlist_path, f"./{self.config.LIST_FILE}")
 
-                lfile_path = hf_hub_download(repo_id=self.repo_ids["log"], filename=self.config.LOG_FILE)
-                shutil.copyfile(lfile_path, f"./{self.config.LOG_FILE}")
+                model_list_path = hf_hub_download(self.repo_ids["model_list"], self.config.LIST_FILE)
+                shutil.copyfile(model_list_path, f"./{self.config.LIST_FILE}")
+
+                log_path = hf_hub_download(self.repo_ids["log"], self.config.LOG_FILE)
+                shutil.copyfile(log_path, f"./{self.config.LOG_FILE}")
 
-                # read logs
-                with open(self.config.LOG_FILE, 'r', encoding='utf-8') as f:
+                with open(self.config.LOG_FILE, "r", encoding="utf-8") as f:
                     lines = f.read().splitlines()
                     old_models = json.loads(lines[0]) if len(lines)>0 else []
                     self.repo_ids["current"] = lines[1] if len(lines)>1 else ""
@@ -313,13 +283,12 @@ class CivitAICrawler:
                 items = r.json().get("items", [])
                 new_ids = [it["id"] for it in items if "id" in it]
 
-                # diff
                 diff_ids = list(set(new_ids) - set(old_models))
                 if diff_ids:
                     mid = diff_ids[0]
                     for attempt in range(1,6):
                         try:
-                            self.process_model(str(mid))  # モデルIDはstrにして渡す
+                            self.process_model(str(mid))
                             break
                         except Exception as e:
                             logger.error(f"[ERR] process_model {mid} (attempt {attempt}): {e}")
@@ -335,7 +304,6 @@ class CivitAICrawler:
 
                     # アップロードログ
                     self.upload_file(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
-                    self.upload_file(self.config.LIST_FILE, self.repo_ids["model_list"], self.config.LIST_FILE)
                 else:
                     with open(self.config.LOG_FILE,'w',encoding='utf-8') as f:
                         f.write(json.dumps(new_ids)+"\n")
@@ -348,6 +316,7 @@ class CivitAICrawler:
                 logger.error(f"[ERR] crawl => {e}")
                 await asyncio.sleep(300)
 
+
 # FastAPI
 config = Config()
 crawler = CivitAICrawler(config)