Upload main.py

main.py

@@ -16,13 +16,10 @@ from fake_useragent import UserAgent
 from fastapi import FastAPI
 from huggingface_hub import HfApi, hf_hub_download, login
 
-# ロギングの設定
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-
 class Config:
-    """設定用のクラス"""
     HUGGINGFACE_API_KEY = os.environ["HUGGINGFACE_API_KEY"]
     CIVITAI_API_TOKEN = os.environ["CIVITAI_API_TOKEN"]
     LOG_FILE = "civitai_backup.log"
@@ -47,10 +44,7 @@ class Config:
         "Content-Type": "application/json"
     }
 
-
 class CivitAICrawler:
-    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス"""
-
     def __init__(self, config: Config):
         import base64
 
@@ -76,14 +70,13 @@ class CivitAICrawler:
         self.setup_routes()
 
     def setup_routes(self):
-        """FastAPIのルーティングを設定する。"""
         @self.app.get("/")
         def read_root():
             now = str(datetime.datetime.now(self.jst))
             description = f"""
-            CivitAIを定期的に周回し新規モデルを {self.repo_ids['current']} にバックアップするspaceです。
-            モデル一覧は https://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.log を参照してください。
-            Status: {now} + currently running :D
+            CivitAIを定期的に周回し新規モデルを {self.repo_ids['current']} にバックアップするSpaceです。
+            モデル一覧は https://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.log 参照。
+            Status: {now} + currently running.
             """
             return description
 
@@ -93,7 +86,6 @@ class CivitAICrawler:
 
     @staticmethod
     def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
-        """Content-Dispositionヘッダーからファイル名を取得する。"""
         if content_disposition:
             parts = content_disposition.split(';')
             for part in parts:
@@ -102,133 +94,120 @@ class CivitAICrawler:
         return default_name
 
     def download_file(self, url: str, destination_folder: str, default_name: str) -> Optional[str]:
-        """指定されたURLからファイルをダウンロードし、指定されたフォルダに保存する。"""
+        """URLからファイルをダウンロードし、destination_folder に保存する。"""
+        os.makedirs(destination_folder, exist_ok=True)  # 念のためフォルダ作成
         try:
-            response = requests.get(url, headers=self.config.HEADERS, stream=True)
-            response.raise_for_status()
+            resp = requests.get(url, headers=self.config.HEADERS, stream=True)
+            resp.raise_for_status()
         except requests.RequestException as e:
             logger.error(f"Failed to download file from {url}: {e}")
             return None
 
-        filename = self.get_filename_from_cd(response.headers.get('content-disposition'), default_name)
+        filename = self.get_filename_from_cd(resp.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)
+        with open(file_path, 'wb') as f:
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
         logger.info(f"Downloaded: {file_path}")
         return file_path
 
     def get_model_info(self, model_id: str) -> dict:
-        """モデルの情報を取得する。"""
         try:
-            response = requests.get(self.config.URLS["modelId"] + str(model_id), headers=self.config.HEADERS)
-            response.raise_for_status()
-            return response.json()
+            resp = requests.get(self.config.URLS["modelId"] + str(model_id), 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 {}
 
     def download_images(self, model_versions: list, folder: str):
-        """画像を images フォルダにまとめてダウンロードする."""
+        """モデル画像を folder/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", []):
+        for ver in model_versions:
+            for img in ver.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)
-            try:
-                resp = requests.get(image_url, stream=True)
-                resp.raise_for_status()
-                with open(local_path, 'wb') as imgf:
-                    for chunk in resp.iter_content(chunk_size=8192):
-                        imgf.write(chunk)
-                logger.info(f"Downloaded image: {local_path}")
-            except requests.RequestException as e:
-                logger.error(f"Failed to download image {image_url}: {e}")
+            self.download_file(image_url, images_folder, image_name)
 
     def save_html_content(self, model_page_url: str, folder: str):
-        """モデルページのHTMLをフォルダ内に保存する."""
+        """model_page_url のHTMLを {folder}/{folder名}.html に保存"""
         try:
             resp = requests.get(model_page_url)
             resp.raise_for_status()
-            html_path = os.path.join(folder, f"{os.path.basename(folder)}.html")
+            html_name = os.path.basename(folder) + ".html"
+            html_path = os.path.join(folder, html_name)
             with open(html_path, 'w', encoding='utf-8') as f:
                 f.write(resp.text)
             logger.info(f"Saved HTML: {html_path}")
         except Exception as e:
-            logger.error(f"Error saving HTML content from {model_page_url}: {e}")
+            logger.error(f"Error saving HTML from {model_page_url}: {e}")
 
     def save_model_info_json(self, model_info: dict, folder: str):
-        """モデル情報(json)の保存"""
+        """model_info.json を folder に保存"""
         info_path = os.path.join(folder, "model_info.json")
         try:
             with open(info_path, 'w', encoding='utf-8') as f:
                 json.dump(model_info, f, indent=2)
             logger.info(f"Saved model_info.json: {info_path}")
         except Exception as e:
-            logger.error(f"Failed to save model info JSON: {e}")
+            logger.error(f"Failed to save model_info.json: {e}")
 
     def download_and_process_versions(self, model_versions: list, folder: str):
         """
-        最新バージョン + 古いバージョンをすべて1つのフォルダにダウンロードして
-        最後にまとめて暗号化アップロードする。
+        最新バージョンは folder/ に、
+        古いバージョンは folder/old_versions/ にまとめて保存。
         """
+        if not model_versions:
+            return
 
-        # 1) 最新バージョンは folder に
-        latest_version = model_versions[0]
-        for file_info in latest_version.get("files", []):
+        # 最新バージョン => folder
+        latest_ver = model_versions[0]
+        for file_info in latest_ver.get("files", []):
             download_url = file_info["downloadUrl"]
             file_name = file_info["name"]
-            local_path = self.download_file(download_url, folder, file_name)
-            # リトライ処理など省略
+            self.download_file(download_url, folder, file_name)
 
-        # 2) 古いバージョンを "folder/old_versions" にまとめる
+        # 古いバージョン => 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:]:
-                for file_info in version.get("files", []):
-                    file_name = file_info["name"]
-                    download_url = file_info["downloadUrl"]
-                    local_path = self.download_file(download_url, old_versions_folder, file_name)
+            oldv_folder = os.path.join(folder, "old_versions")
+            os.makedirs(oldv_folder, exist_ok=True)
+            for v in model_versions[1:]:
+                for f_info in v.get("files", []):
+                    dl_url = f_info["downloadUrl"]
+                    f_name = f_info["name"]
+                    self.download_file(dl_url, oldv_folder, f_name)
 
     def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
         """
-        1. rclone copy local_folder => cryptLocal:   (フォルダ名はrcloneが勝手に暗号化生成)
-        2. 差分検知で "./encrypted" に作られた暗号フォルダ名を取得
-        3. そのフォルダ名を path_in_repo として Hugging Face にアップロード
-        4. ローカル(平文+暗号)フォルダ削除
-        5. 戻り値は "実際に作られた暗号フォルダ名"
+        rclone copy local_folder => cryptLocal:
+        => 差分検知で "encrypted/xxxxxx" を発見 -> upload_folder -> 削除
+        => 戻り値は暗号フォルダ名
         """
         if not os.path.isdir(local_folder):
-            logger.error(f"[encrypt_and_upload_folder] Not a directory: {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)
 
-        # 既存の暗号フォルダを削除
         before_set = set(os.listdir(encrypted_base_dir))
-        for item in before_set:
-            item_path = os.path.join(encrypted_base_dir, item)
+        # cleanup old stuff
+        for itm in before_set:
+            itm_path = os.path.join(encrypted_base_dir, itm)
             try:
-                if os.path.isfile(item_path):
-                    os.remove(item_path)
+                if os.path.isfile(itm_path):
+                    os.remove(itm_path)
                 else:
-                    shutil.rmtree(item_path)
-                logger.info(f"[CLEANUP] Removed old encrypted item: {item_path}")
+                    shutil.rmtree(itm_path)
+                logger.info(f"[CLEANUP] Removed {itm_path}")
             except Exception as e:
-                logger.warning(f"[CLEANUP] Failed to remove {item_path}: {e}")
+                logger.warning(f"[CLEANUP] Failed to remove {itm_path}: {e}")
 
-        # === (1) rclone copy local_folder => cryptLocal: ===
-        #     これにより /home/user/app/encrypted/<ランダム暗号フォルダ> が作成される
+        # rclone copy local_folder => cryptLocal:
         try:
             subprocess.run(
                 ["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"],
@@ -239,39 +218,35 @@ class CivitAICrawler:
             logger.error(f"rclone copy failed: {e}")
             return None
 
-        # === (2) 差分検知: 新しく生成された暗号フォルダ名を取得 ===
         after_set = set(os.listdir(encrypted_base_dir))
         diff = after_set - before_set
         if not diff:
-            logger.error("[ERROR] No new directory appeared in ./encrypted after rclone copy.")
+            logger.error("[ERROR] No new directory in ./encrypted after rclone copy.")
             return None
         if len(diff) > 1:
-            logger.warning(f"[WARN] Multiple new directories created: {diff}")
+            logger.warning(f"[WARN] multiple new dirs? {diff}")
 
-        enc_folder_name = diff.pop()  # 1個だけ取り出す
+        enc_folder_name = diff.pop()
         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
 
-        # === (3) Hugging Face にアップロード ===
-        #     path_in_repo も "enc_folder_name" をそのまま使う
+        # upload_folder
         try:
             self.upload_folder(enc_folder_path, path_in_repo=enc_folder_name)
-            logger.info(f"[OK] Uploaded encrypted folder: {enc_folder_path}")
+            logger.info(f"[OK] Uploaded {enc_folder_path}")
         except Exception as e:
-            logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")
+            logger.error(f"Failed to upload {enc_folder_path}: {e}")
 
-        # === (4) ローカル削除 (平文フォルダ + 暗号化フォルダ)
+        # cleanup local
         try:
             shutil.rmtree(local_folder)
             shutil.rmtree(enc_folder_path)
-            logger.info(f"[CLEANUP] Removed local folder: {local_folder} & {enc_folder_path}")
+            logger.info(f"[CLEANUP] Removed {local_folder} & {enc_folder_path}")
         except Exception as e:
             logger.warning(f"[CLEANUP] Could not remove local folders: {e}")
 
-        # === (5) 実際の暗号フォルダ名を返す
         return enc_folder_name
 
     def upload_file(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None):
@@ -279,7 +254,7 @@ class CivitAICrawler:
             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:
@@ -301,7 +276,7 @@ class CivitAICrawler:
                     attempt = 0
                     continue
                 elif "you can retry this action in about 1 hour" in error_message:
-                    logger.warning("Rate limit hit. Waiting 1 hour...")
+                    logger.warning("Rate limit. Wait 1hr.")
                     time.sleep(3600)
                     attempt -= 1
                 else:
@@ -324,24 +299,24 @@ class CivitAICrawler:
                     repo_id=self.repo_ids['current'],
                     path_in_repo=path_in_repo
                 )
-                logger.info(f"Uploaded folder: {folder_path} to {self.repo_ids['current']} at {path_in_repo}")
+                logger.info(f"Uploaded folder: {folder_path} => {self.repo_ids['current']}:{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.")
+                    logger.warning("File limit exceeded, creating 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...")
+                    logger.warning("Rate limit. Waiting 1hr.")
                     time.sleep(3600)
                     attempt -= 1
                 else:
                     if attempt < max_retries:
-                        logger.warning(f"Failed to upload folder {folder_path}, retry {attempt}/{max_retries}")
+                        logger.warning(f"Failed to upload folder {folder_path}, attempt {attempt}/{max_retries}")
                     else:
                         logger.error(f"Failed after {max_retries} attempts: {e}")
                         raise
@@ -381,7 +356,6 @@ class CivitAICrawler:
             return []
 
     def process_model(self, model_url: str):
-        """1つのモデルをフォルダにまとめてダウンロード→暗号化→アップロード"""
         try:
             model_id = model_url.rstrip("/").split("/")[-1]
             model_info = self.get_model_info(model_id)
@@ -399,76 +373,76 @@ class CivitAICrawler:
             folder_name += "_" + str(uuid.uuid4())[:8]
             os.makedirs(folder_name, exist_ok=True)
 
-            # 最新 + 古いバージョンすべて 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)
+            # 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_model_info_json(model_info, folder_name)
 
-            # フォルダごと暗号化
-            enc_subfolder = self.encrypt_and_upload_folder(folder_name)
-            if enc_subfolder is None:
-                enc_subfolder = "[ENCRYPT_FAILED]"
+            # 最後にフォルダごとアップ
+            enc_folder = self.encrypt_and_upload_folder(folder_name)
+            if enc_folder is None:
+                enc_folder = "[ENCRYPT_FAILED]"
 
-            hf_enc_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{enc_subfolder}"
+            hf_enc_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{enc_folder}"
             with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
-                f.write(f"{model_info.get('name', 'UnnamedModel')} (ID:{model_id}): {hf_enc_url}\n")
+                f.write(f"{model_info.get('name','Unknown')} (ID:{model_id}): {hf_enc_url}\n")
 
         except Exception as e:
-            logger.error(f"Error in process_model ({model_url}): {e}")
+            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_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}")
 
-                # ログを読み込み
                 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 ""
 
-                # 新着モデル確認
                 r = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                 r.raise_for_status()
                 latest_models = r.json().get("items", [])
-                latest_model_ids = [m["id"] for m in latest_models if "id" in m]
+                latest_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 model IDs found: {new_models}")
-                    model_id = new_models[0]
+                new_ids = list(set(latest_ids) - set(old_models))
+                if new_ids:
+                    logger.info(f"New model IDs found: {new_ids}")
+                    mid = new_ids[0]
 
-                    for attempt in range(1, 6):
+                    for attempt in range(1,6):
                         try:
-                            self.process_model(self.config.URLS["modelId"] + str(model_id))
+                            self.process_model(f"{self.config.URLS['modelId']}{mid}")
                             break
                         except Exception as e:
-                            logger.error(f"Failed to process model {model_id} (attempt {attempt}/5): {e}")
+                            logger.error(f"Failed model {mid} (attempt {attempt}/5): {e}")
                             if attempt == 5:
-                                logger.error(f"Skipping model {model_id} after 5 failures.")
+                                logger.error(f"Skipping model {mid}")
                             else:
                                 await asyncio.sleep(2)
 
-                    old_models.append(model_id)
+                    old_models.append(mid)
                     with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
-                        f.write(json.dumps(old_models) + "\n")
+                        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}")
+                    logger.info(f"Updated log with new model ID: {mid}")
 
                     self.upload_file(
                         file_path=self.config.LOG_FILE,
@@ -482,9 +456,9 @@ class CivitAICrawler:
                     )
                 else:
                     with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
-                        f.write(json.dumps(latest_model_ids) + "\n")
+                        f.write(json.dumps(latest_ids)+"\n")
                         f.write(f"{self.repo_ids['current']}\n")
-                    logger.info(f"No new models. Updated log: {self.config.LOG_FILE}")
+                    logger.info("No new models found. Updated log.")
                     self.upload_file(
                         file_path=self.config.LOG_FILE,
                         repo_id=self.repo_ids["log"],
@@ -493,13 +467,11 @@ class CivitAICrawler:
                     logger.info("Uploaded log file.")
                     await asyncio.sleep(60)
                     continue
-
             except Exception as e:
                 logger.error(f"Error in crawl loop: {e}")
                 await asyncio.sleep(300)
 
-
-# FastAPIアプリケーション
+# FastAPI
 config = Config()
 crawler = CivitAICrawler(config)
 app = crawler.app