Update main.py

main.py

@@ -1,142 +1,290 @@
+import asyncio
+import datetime
+import json
+import logging
 import os
-import subprocess
+import re
 import shutil
+import subprocess
+import time
 import uuid
-import base64
-import requests
+from typing import Optional
 
-from huggingface_hub import HfApi, login
+import requests
+from bs4 import BeautifulSoup
+from fake_useragent import UserAgent
+from fastapi import FastAPI
+from huggingface_hub import HfApi, hf_hub_download, login
 
-# ログ出力用
-import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# === 環境変数 ===
-HUGGINGFACE_API_KEY = os.environ.get("HUGGINGFACE_API_KEY", "")
-RCLONE_CONF_BASE64 = os.environ.get("RCLONE_CONF_BASE64", "")
-REPO_ID = os.environ.get("REPO_ID", "username/testrepo")  
-# ↑ あなたのアップロード先リポジトリ (例: "ttttdiva/CivitAI_Auto12")
-
-def setup_rclone_conf():
-    """RCLONE_CONF_BASE64 から rclone.conf を復元"""
-    if not RCLONE_CONF_BASE64:
-        logger.warning("[WARN] RCLONE_CONF_BASE64 not set. 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_CONF_BASE64))
-    os.environ["RCLONE_CONFIG"] = conf_path
-    logger.info(f"[INFO] rclone.conf created => {conf_path}")
-
-def download_file(url: str, dest_folder: str, filename: str):
-    """URL をダウンロードして dest_folder/filename に保存"""
-    os.makedirs(dest_folder, exist_ok=True)
-    try:
-        r = requests.get(url, stream=True)
-        r.raise_for_status()
-        filepath = os.path.join(dest_folder, filename)
-        with open(filepath, 'wb') as f:
-            for chunk in r.iter_content(chunk_size=8192):
+class Config:
+    HUGGINGFACE_API_KEY = os.environ["HUGGINGFACE_API_KEY"]
+    CIVITAI_API_TOKEN = os.environ["CIVITAI_API_TOKEN"]
+    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": ""
+    }
+    URLS = {
+        "latest": "https://civitai.com/api/v1/models?sort=Newest",
+        "modelPage": "https://civitai.com/models/",
+        "modelId": "https://civitai.com/api/v1/models/",
+        "modelVersionId": "https://civitai.com/api/v1/model-versions/",
+        "hash": "https://civitai.com/api/v1/model-versions/by-hash/"
+    }
+    JST = datetime.timezone(datetime.timedelta(hours=9))
+    UA = UserAgent()
+    HEADERS = {
+        'Authorization': f'Bearer {CIVITAI_API_TOKEN}',
+        'User-Agent': UA.random,
+        "Content-Type": "application/json"
+    }
+
+class CivitAICrawler:
+    def __init__(self, config: Config):
+        self.config = config
+        self.api = HfApi()
+        self.app = FastAPI()
+        self.repo_ids = self.config.REPO_IDS.copy()
+        self.jst = self.config.JST
+        self.setup_rclone_conf()
+        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}")
+
+    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"
+            }
+
+        @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)
+        try:
+            resp = requests.get(url, headers=self.config.HEADERS, stream=True)
+            resp.raise_for_status()
+        except requests.RequestException as e:
+            logger.error(f"[ERR] download_file fail: {e}")
+            return None
+
+        file_path = os.path.join(destination_folder, default_name)
+        with open(file_path, 'wb') as f:
+            for chunk in resp.iter_content(chunk_size=8192):
                 f.write(chunk)
-        logger.info(f"[OK] Downloaded: {filepath}")
-        return filepath
-    except Exception as e:
-        logger.error(f"[ERR] download_file failed: {e}")
-        return None
-
-def encrypt_and_upload_folder(local_folder: str):
-    """
-    1) rclone copy local_folder => cryptLocal:
-       => /home/user/app/encrypted/<暗号フォルダ>
-    2) その暗号フォルダを Hugging Face へアップロード
-    3) ローカル削除
-    """
-    if not os.path.isdir(local_folder):
-        logger.error(f"[ERR] {local_folder} is not a directory.")
-        return
-
-    encrypted_dir = os.path.join(os.getcwd(), "encrypted")
-    os.makedirs(encrypted_dir, exist_ok=True)
-
-    # 差分検知
-    before = set(os.listdir(encrypted_dir))
-
-    # rclone copy
-    try:
-        subprocess.run(["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"], check=True)
-        logger.info(f"[OK] rclone copy {local_folder} => cryptLocal:")
-    except subprocess.CalledProcessError as e:
-        logger.error(f"[ERR] rclone copy failed: {e}")
-        return
-
-    after = set(os.listdir(encrypted_dir))
-    diff = after - before
-    if not diff:
-        logger.error("[ERR] No new folder in ./encrypted after rclone copy.")
-        return
-    if len(diff) > 1:
-        logger.warning(f"[WARN] multiple new folders? {diff}")
-    enc_folder_name = diff.pop()
-    enc_folder_path = os.path.join(encrypted_dir, enc_folder_name)
-    logger.info(f"[DEBUG] enc_folder_path => {enc_folder_path}")
-
-    if not os.path.isdir(enc_folder_path):
-        logger.error(f"[ERR] {enc_folder_path} is not a directory.")
-        return
-
-    # Hugging Face にアップロード (フォルダ)
-    try:
-        api = HfApi()
-        # subfolder_label = enc_folder_name  # そのまま
-        subfolder_label = enc_folder_name  # そのまま
-
-        api.upload_folder(
-            folder_path=enc_folder_path,
-            repo_id=REPO_ID,
-            path_in_repo=subfolder_label
-        )
-        logger.info(f"[OK] uploaded folder => {enc_folder_path} to {REPO_ID}:{subfolder_label}")
-    except Exception as e:
-        logger.error(f"[ERR] upload_folder failed: {e}")
-
-    # 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}")
-
-def main():
-    logger.info("===== Starting minimal test =====")
-
-    # 1) rclone.conf 復元
-    setup_rclone_conf()
-
-    # 2) Hugging Face login
-    if not HUGGINGFACE_API_KEY:
-        logger.error("[ERR] HUGGINGFACE_API_KEY not set.")
-        return
-    login(token=HUGGINGFACE_API_KEY, add_to_git_credential=True)
-    logger.info("[OK] HF login success")
-
-    # 3) ローカルにフォルダ + ファイル
-    test_folder = "MyTestModel"
-    if os.path.exists(test_folder):
-        shutil.rmtree(test_folder)
-    os.makedirs(test_folder, exist_ok=True)
-
-    # ダウンロード例 (画像2つ + info.json 1つ)
-    download_file("https://picsum.photos/200/300", test_folder, "image1.jpg")
-    download_file("https://picsum.photos/300/300", test_folder, "image2.jpg")
-
-    info_path = os.path.join(test_folder, "model_info.json")
-    with open(info_path, 'w', encoding='utf-8') as f:
-        f.write('{"model":"test","desc":"some description"}')
-    logger.info(f"[OK] Created {info_path}")
-
-    # 4) 暗号化してアップロード
-    encrypt_and_upload_folder(test_folder)
-
-if __name__ == "__main__":
-    main()
+        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 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()
+        except Exception as e:
+            logger.error(f"[ERR] get_model_info({model_id}): {e}")
+            return {}
+
+    def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
+        """
+        rclone copy local_folder => cryptLocal:
+        => diff => 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))
+
+        try:
+            subprocess.run(["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"], check=True)
+            logger.info(f"[OK] rclone copy {local_folder} => cryptLocal:")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"[ERR] rclone copy failed: {e}")
+            return None
+
+        after = set(os.listdir(encrypted_dir))
+        diff = after - before
+        if not diff:
+            logger.error("[ERR] No new dir in ./encrypted after rclone 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.")
+            return None
+
+        # upload_folder to HF
+        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}")
+        except Exception as e:
+            logger.error(f"[ERR] HF upload_folder fail: {e}")
+
+        # cleanup local
+        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_id: str):
+        # ダウンロード + 暗号化 + アップ
+        try:
+            minfo = self.get_model_info(model_id)
+            if not minfo or "modelVersions" not in minfo:
+                logger.error(f"[ERR] No modelVersions for {model_id}")
+                return
+
+            versions = minfo["modelVersions"]
+            # local folder
+            folder_name = re.sub(r'[\\/*?:"<>|]', '_', minfo.get("name", "UnknownModel"))
+            folder_name += f"_{uuid.uuid4()[:8]}"
+            os.makedirs(folder_name, exist_ok=True)
+
+            # ダウンロード (最新/古い) => まとめて folder_name
+            # 画像 => folder_name/images
+            # etc...
+            self.download_versions(versions, folder_name)
+            self.download_images(versions, folder_name)
+            # ここでHTMLや model_info.json 保存もやるなら追加:
+            # ...
+            enc_f = self.encrypt_and_upload_folder(folder_name)
+            # => HF upload + cleanup
+
+        except Exception as e:
+            logger.error(f"[ERR] process_model({model_id}) => {e}")
+
+    def download_versions(self, model_versions: list, folder: str):
+        # すべて folder/ にまとめる or old_versions subfolder
+        # 例: 最新 => folder, old => 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)
+            for v in model_versions[1:]:
+                for f_info in v.get("files", []):
+                    url = f_info["downloadUrl"]
+                    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)
+
+    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)
+                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}")
+
+                # read logs
+                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 ""
+
+                # get newest
+                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 = list(set(ids)-set(old_models))
+                if new_ids:
+                    mid = new_ids[0]
+                    for attempt in range(1,6):
+                        try:
+                            self.process_model(str(mid))
+                            break
+                        except Exception as e:
+                            logger.error(f"[ERR] process_model {mid}, attempt {attempt}: {e}")
+                            if attempt==5:
+                                logger.error("Skipping model 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:
+                        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")
+                        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")
+                    await asyncio.sleep(60)
+                    continue
+            except Exception as e:
+                logger.error(f"[ERR] crawl loop => {e}")
+                await asyncio.sleep(300)
+
+# === FastAPI
+config = Config()
+crawler = CivitAICrawler(config)
+app = crawler.app