Update main.py

main.py

@@ -25,9 +25,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",
@@ -55,83 +55,99 @@ class CivitAICrawler:
         self.setup_routes()
 
     def setup_rclone_conf(self):
-        # RCLONE_CONF_BASE64→rclone.conf復元
         import base64
         rclone_b64 = os.environ.get("RCLONE_CONF_BASE64", "")
-        if not rclone_b64:
-            logger.warning("[WARN] RCLONE_CONF_BASE64 is empty, rclone may fail.")
-            return
-        conf_dir = ".rclone_config"
-        os.makedirs(conf_dir, exist_ok=True)
-        conf_path = os.path.join(conf_dir, "rclone.conf")
-        with open(conf_path, "wb") as f:
-            f.write(base64.b64decode(rclone_b64))
-        os.environ["RCLONE_CONFIG"] = conf_path
-        logger.info(f"[OK] rclone.conf => {conf_path}")
+        if rclone_b64:
+            conf_dir = ".rclone_config"
+            os.makedirs(conf_dir, exist_ok=True)
+            conf_path = os.path.join(conf_dir, "rclone.conf")
+            with open(conf_path, "wb") as f:
+                f.write(base64.b64decode(rclone_b64))
+            os.environ["RCLONE_CONFIG"] = conf_path
+            logger.info(f"[OK] Created rclone.conf => {conf_path}")
+        else:
+            logger.warning("[WARN] RCLONE_CONF_BASE64 is empty. rclone may fail.")
 
     def setup_routes(self):
         @self.app.get("/")
         def read_root():
             now = str(datetime.datetime.now(self.jst))
             return {
-                "description": f"CivitAI crawler. Current time: {now}",
-                "repo_current": self.repo_ids["current"],
-                "note": "Startup event => self.crawl() loop"
+                "description": f"CivitAI crawler. Time: {now}",
+                "repo_current": self.repo_ids["current"]
             }
 
         @self.app.on_event("startup")
         async def startup_event():
             asyncio.create_task(self.crawl())
 
-    def download_file(self, url: str, destination_folder: str, default_name: str) -> Optional[str]:
-        os.makedirs(destination_folder, exist_ok=True)
+    # ここで「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:
-            resp = requests.get(url, headers=self.config.HEADERS, stream=True)
-            resp.raise_for_status()
+            r = requests.get(url, headers=self.config.HEADERS, stream=True)
+            r.raise_for_status()
         except requests.RequestException as e:
-            logger.error(f"[ERR] download_file fail: {e}")
+            logger.error(f"[ERR] download_file: {e}")
             return None
 
-        file_path = os.path.join(destination_folder, default_name)
+        file_path = os.path.join(dest_folder, filename)
         with open(file_path, 'wb') as f:
-            for chunk in resp.iter_content(chunk_size=8192):
+            for chunk in r.iter_content(chunk_size=8192):
                 f.write(chunk)
-        logger.info(f"[OK] Downloaded: {file_path}")
+        logger.info(f"[OK] Downloaded => {file_path}")
         return file_path
 
-    def get_filename_from_cd(self, cd: Optional[str], default_name: str) -> str:
-        if cd:
-            parts = cd.split(';')
-            for p in parts:
-                if "filename=" in p:
-                    return p.split("=")[1].strip().strip('"')
-        return default_name
+    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)
 
-    def get_model_info(self, model_id: str) -> dict:
         try:
-            r = requests.get(f"{self.config.URLS['modelId']}{model_id}", headers=self.config.HEADERS)
-            r.raise_for_status()
-            return r.json()
+            self.api.upload_file(
+                path_or_fileobj=file_path,
+                repo_id=repo_id,
+                path_in_repo=path_in_repo
+            )
+            logger.info(f"[OK] Uploaded file => {repo_id}:{path_in_repo}")
+        except Exception as e:
+            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)
+
+        try:
+            self.api.upload_folder(
+                folder_path=folder_path,
+                repo_id=self.repo_ids["current"],
+                path_in_repo=path_in_repo
+            )
+            logger.info(f"[OK] uploaded folder => {folder_path} => {self.repo_ids['current']}:{path_in_repo}")
         except Exception as e:
-            logger.error(f"[ERR] get_model_info({model_id}): {e}")
-            return {}
+            logger.error(f"[ERR] upload_folder: {e}")
 
     def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
-        """
-        rclone copy local_folder => cryptLocal:
-        => diff => upload_folder => 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
+
         encrypted_dir = os.path.join(os.getcwd(), "encrypted")
         os.makedirs(encrypted_dir, exist_ok=True)
 
         before = set(os.listdir(encrypted_dir))
-
+        # rclone copy
+        cmd = ["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"]
+        logger.info(f"[CMD] {' '.join(cmd)}")
         try:
-            subprocess.run(["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"], check=True)
-            logger.info(f"[OK] rclone copy {local_folder} => cryptLocal:")
+            subprocess.run(cmd, check=True)
+            logger.info(f"[OK] rclone copy => cryptLocal:")
         except subprocess.CalledProcessError as e:
             logger.error(f"[ERR] rclone copy failed: {e}")
             return None
@@ -139,85 +155,47 @@ class CivitAICrawler:
         after = set(os.listdir(encrypted_dir))
         diff = after - before
         if not diff:
-            logger.error("[ERR] No new dir in ./encrypted after rclone copy.")
+            logger.error("[ERR] no new directory in ./encrypted after copy")
             return None
         if len(diff) > 1:
             logger.warning(f"[WARN] multiple new dirs => {diff}")
-        enc_folder_name = diff.pop()
-        enc_folder_path = os.path.join(encrypted_dir, enc_folder_name)
-
-        if not os.path.isdir(enc_folder_path):
-            logger.error(f"[ERR] {enc_folder_path} is not a directory.")
+        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
 
-        # upload_folder to HF
+        # HF upload folder
         try:
-            self.api.upload_folder(
-                folder_path=enc_folder_path,
-                repo_id=self.repo_ids["current"],
-                path_in_repo=enc_folder_name
-            )
-            logger.info(f"[OK] uploaded {enc_folder_path} => {self.repo_ids['current']}:{enc_folder_name}")
+            self.upload_folder(enc_path, path_in_repo=enc_name)
         except Exception as e:
-            logger.error(f"[ERR] HF upload_folder fail: {e}")
+            logger.error(f"[ERR] encrypt_and_upload_folder => upload_folder: {e}")
 
-        # cleanup local
+        # cleanup
         shutil.rmtree(local_folder, ignore_errors=True)
-        shutil.rmtree(enc_folder_path, ignore_errors=True)
-        logger.info(f"[CLEANUP] removed {local_folder} & {enc_folder_path}")
-
-        return enc_folder_name
-
-    def process_model(self, model_url: str):
-        try:
-            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
-
-            # 1) 単にモデル名をサニタイズしたディレクトリにまとめる
-            folder_name = re.sub(r'[\\/*?:"<>|]', '_', model_info.get("name", "UnknownModel"))
-
-            # フォルダが既にあれば削除し、新規に作り直す (上書き)
-            if os.path.exists(folder_name):
-                shutil.rmtree(folder_name)
-            os.makedirs(folder_name, exist_ok=True)
-
-            # ダウンロード (最新/old_versions), 画像, HTML, model_info.json など
-            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)
-
-            # 暗号化アップロード
-            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}"
-            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")
-
-        except Exception as e:
-            logger.error(f"Error in process_model({model_url}): {e}")
-
-    def download_versions(self, model_versions: list, folder: str):
-        # すべて folder/ にまとめる or old_versions subfolder
-        # 例: 最新 => folder, old => folder/old_versions
+        shutil.rmtree(enc_path, ignore_errors=True)
+        logger.info(f"[CLEANUP] removed {local_folder} & {enc_path}")
+        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)
@@ -227,31 +205,67 @@ class CivitAICrawler:
                     fname = f_info["name"]
                     self.download_file(url, ov_folder, fname)
 
-    def download_images(self, model_versions: list, folder: str):
-        images_folder = os.path.join(folder, "images")
-        os.makedirs(images_folder, exist_ok=True)
-        for v in model_versions:
-            for im in v.get("images", []):
-                iurl = im["url"]
-                iname = os.path.basename(iurl)
-                self.download_file(iurl, images_folder, iname)
+    def get_model_info_json(self, model_info: dict, folder: str):
+        """model_info.jsonを保存"""
+        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"[OK] saved model_info.json => {info_path}")
+        except Exception as e:
+            logger.error(f"[ERR] saving model_info.json: {e}")
+
+    def process_model(self, model_id: str):
+        """本体処理: フォルダにまとめてDL => encrypt => upload"""
+        try:
+            info = self.get_model_info(model_id)
+            if not info or "modelVersions" not in info:
+                logger.error(f"[ERR] no modelVersions for {model_id}")
+                return
+            versions = info["modelVersions"]
+            folder_name = re.sub(r'[\\/*?:"<>|]', '_', info.get("name","UnknownModel"))
+            # local_folderベースに作りたいならこう↓
+            # local_folder = "local_folder"
+            # os.makedirs(local_folder, exist_ok=True)
+            # folder_path = os.path.join(local_folder, folder_name)
+            # ここでは、とりあえずapp直下に folder_name作成
+            folder_path = folder_name
+
+            if os.path.exists(folder_path):
+                shutil.rmtree(folder_path)
+            os.makedirs(folder_path, exist_ok=True)
+            logger.info(f"[OK] created folder => {folder_path}")
+
+            # ダウンロード
+            self.download_and_process_versions(versions, folder_path)
+
+            # 画像
+            self.download_images(versions, folder_path)
+
+            # model_info.json
+            self.get_model_info_json(info, folder_path)
+
+            # 暗号化アップ
+            enc_folder = self.encrypt_and_upload_folder(folder_path)
+            if enc_folder is None:
+                enc_folder = "[ENCRYPT_FAILED]"
+            # (Optional)  model_list.log に追記 など
+        except Exception as e:
+            logger.error(f"Error in process_model({model_id}): {e}")
 
     async def crawl(self):
         while True:
             try:
-                # HF Login
                 login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
-
-                # model_list.log
-                mlist_path = hf_hub_download(self.repo_ids["model_list"], self.config.LIST_FILE)
+                # 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}")
 
-                # log_file
-                log_path = hf_hub_download(self.repo_ids["log"], self.config.LOG_FILE)
-                shutil.copyfile(log_path, f"./{self.config.LOG_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}")
 
                 # 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 ""
@@ -260,43 +274,44 @@ class CivitAICrawler:
                 r = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                 r.raise_for_status()
                 items = r.json().get("items", [])
-                ids = [it["id"] for it in items if "id" in it]
+                new_ids = [it["id"] for it in items if "id" in it]
 
-                new_ids = list(set(ids)-set(old_models))
-                if new_ids:
-                    mid = new_ids[0]
+                # 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))
+                            self.process_model(str(mid))  # モデルIDはstrにして渡す
                             break
                         except Exception as e:
-                            logger.error(f"[ERR] process_model {mid}, attempt {attempt}: {e}")
+                            logger.error(f"[ERR] process_model {mid} (attempt {attempt}): {e}")
                             if attempt==5:
-                                logger.error("Skipping model after 5 fails")
+                                logger.error(f"[SKIP] model {mid} after 5 fails")
                             else:
                                 await asyncio.sleep(2)
-                    # update logs
+
                     old_models.append(mid)
-                    with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
+                    with open(self.config.LOG_FILE,'w',encoding='utf-8') as f:
                         f.write(json.dumps(old_models)+"\n")
                         f.write(self.repo_ids["current"]+"\n")
 
+                    # アップロードログ
                     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:
-                    # no new
-                    with open(self.config.LOG_FILE,"w",encoding="utf-8") as f:
-                        f.write(json.dumps(ids)+"\n")
+                    with open(self.config.LOG_FILE,'w',encoding='utf-8') as f:
+                        f.write(json.dumps(new_ids)+"\n")
                         f.write(self.repo_ids["current"]+"\n")
                     self.upload_file(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
-                    logger.info("No new models => wait 60s")
+                    logger.info("[INFO] no new models => sleep(60)")
                     await asyncio.sleep(60)
                     continue
             except Exception as e:
-                logger.error(f"[ERR] crawl loop => {e}")
+                logger.error(f"[ERR] crawl => {e}")
                 await asyncio.sleep(300)
 
-# === FastAPI
+# FastAPI
 config = Config()
 crawler = CivitAICrawler(config)
 app = crawler.app