Update main.py

main.py

@@ -1,477 +1,142 @@
-import asyncio
-import datetime
-import json
-import logging
 import os
-import re
-import shutil
 import subprocess
-import time
+import shutil
 import uuid
-from typing import Optional
-
+import base64
 import requests
-from bs4 import BeautifulSoup
-from fake_useragent import UserAgent
-from fastapi import FastAPI
-from huggingface_hub import HfApi, hf_hub_download, login
 
+from huggingface_hub import HfApi, login
+
+# ログ出力用
+import logging
 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"
-    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):
-        import base64
-
-        rclone_conf_base64 = os.environ.get("RCLONE_CONF_BASE64")
-        if rclone_conf_base64:
-            config_dir = os.path.join(os.getcwd(), ".rclone_config")
-            os.makedirs(config_dir, exist_ok=True)
-
-            conf_path = os.path.join(config_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] Created rclone.conf at {conf_path}")
-        else:
-            logger.warning("[WARN] RCLONE_CONF_BASE64 not found; rclone may fail.")
-
-        self.config = config
-        self.api = HfApi()
-        self.app = FastAPI()
-        self.repo_ids = self.config.REPO_IDS.copy()
-        self.jst = self.config.JST
-        self.setup_routes()
-
-    def setup_routes(self):
-        @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.
-            """
-            return description
-
-        @self.app.on_event("startup")
-        async def startup_event():
-            asyncio.create_task(self.crawl())
-
-    @staticmethod
-    def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
-        if content_disposition:
-            parts = content_disposition.split(';')
-            for part in parts:
-                if "filename=" in part:
-                    return part.split("=")[1].strip().strip('"')
-        return default_name
-
-    def download_file(self, url: str, destination_folder: str, default_name: str) -> Optional[str]:
-        """URLからファイルをダウンロードし、destination_folder に保存する。"""
-        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"Failed to download file from {url}: {e}")
-            return None
-
-        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 f:
-            for chunk in resp.iter_content(chunk_size=8192):
+# === 環境変数 ===
+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):
                 f.write(chunk)
-        logger.info(f"Downloaded: {file_path}")
-        return file_path
-
-    def get_model_info(self, model_id: str) -> dict:
-        try:
-            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):
-        """モデル画像を folder/images にダウンロード"""
-        images_folder = os.path.join(folder, "images")
-        os.makedirs(images_folder, exist_ok=True)
-        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)
-            self.download_file(image_url, images_folder, image_name)
-
-    def save_html_content(self, model_page_url: str, folder: str):
-        """model_page_url のHTMLを {folder}/{folder名}.html に保存"""
-        try:
-            resp = requests.get(model_page_url)
-            resp.raise_for_status()
-            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 from {model_page_url}: {e}")
-
-    def save_model_info_json(self, model_info: dict, folder: str):
-        """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}")
-
-    def download_and_process_versions(self, model_versions: list, folder: str):
-        """
-        最新バージョンは folder/ に、
-        古いバージョンは folder/old_versions/ にまとめて保存。
-        """
-        if not model_versions:
-            return
-
-        # 最新バージョン => folder
-        latest_ver = model_versions[0]
-        for file_info in latest_ver.get("files", []):
-            download_url = file_info["downloadUrl"]
-            file_name = file_info["name"]
-            self.download_file(download_url, folder, file_name)
-
-        # 古いバージョン => folder/old_versions
-        if len(model_versions) > 1:
-            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]:
-        """
-        rclone copy local_folder => cryptLocal:
-        => 差分検知で "encrypted/xxxxxx" を発見 -> upload_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)
-
-        before_set = set(os.listdir(encrypted_base_dir))
-        # cleanup old stuff
-        for itm in before_set:
-            itm_path = os.path.join(encrypted_base_dir, itm)
-            try:
-                if os.path.isfile(itm_path):
-                    os.remove(itm_path)
-                else:
-                    shutil.rmtree(itm_path)
-                logger.info(f"[CLEANUP] Removed {itm_path}")
-            except Exception as e:
-                logger.warning(f"[CLEANUP] Failed to remove {itm_path}: {e}")
-
-        # rclone copy local_folder => cryptLocal:
-        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"rclone copy failed: {e}")
-            return None
-
-        after_set = set(os.listdir(encrypted_base_dir))
-        diff = after_set - before_set
-        if not diff:
-            logger.error("[ERROR] No new directory 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_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
-
-        # upload_folder
-        try:
-            self.upload_folder(enc_folder_path, path_in_repo=enc_folder_name)
-            logger.info(f"[OK] Uploaded {enc_folder_path}")
-        except Exception as e:
-            logger.error(f"Failed to upload {enc_folder_path}: {e}")
-
-        # cleanup local
-        try:
-            shutil.rmtree(local_folder)
-            shutil.rmtree(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}")
-
-        return enc_folder_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)
-
-        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. Wait 1hr.")
-                    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):
-        if path_in_repo is None:
-            path_in_repo = os.path.basename(folder_path)
-
-        max_retries = 5
-        attempt = 0
-        while attempt < max_retries:
-            try:
-                self.api.upload_folder(
-                    folder_path=folder_path,
-                    repo_id=self.repo_ids['current'],
-                    path_in_repo=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 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. Waiting 1hr.")
-                    time.sleep(3600)
-                    attempt -= 1
-                else:
-                    if 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
-
-    @staticmethod
-    def increment_repo_name(repo_id: str) -> str:
-        match = re.search(r'(\d+)$', repo_id)
-        if match:
-            number = int(match.group(1)) + 1
-            new_repo_id = re.sub(r'\d+$', str(number), repo_id)
-        else:
-            new_repo_id = f"{repo_id}1"
-        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:
-                for line in f:
-                    line = line.strip()
-                    if line:
-                        parts = line.split(": ", 1)
-                        if len(parts) == 2:
-                            modelpage_name, model_hf_url = parts
-                            model_list[model_hf_url] = modelpage_name
-        except Exception as e:
-            logger.error(f"Failed to read model list: {e}")
-        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]
-            return file_paths
-        except Exception as e:
-            logger.error(f"Failed to get repo info for {repo_id}: {e}")
-            return []
-
-    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
-
-            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)
-
-            # ダウンロード(最新+古い)
-            self.download_and_process_versions(versions, folder_name)
-            # 画像
-            self.download_images(versions, 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_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_folder}"
-            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
-                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}")
-
-    async def crawl(self):
-        while True:
-            try:
-                login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
-
-                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_ids = [m["id"] for m in latest_models if "id" in m]
-
-                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):
-                        try:
-                            self.process_model(f"{self.config.URLS['modelId']}{mid}")
-                            break
-                        except Exception as e:
-                            logger.error(f"Failed model {mid} (attempt {attempt}/5): {e}")
-                            if attempt == 5:
-                                logger.error(f"Skipping model {mid}")
-                            else:
-                                await asyncio.sleep(2)
-
-                    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(f"{self.repo_ids['current']}\n")
-                    logger.info(f"Updated log with new model ID: {mid}")
-
-                    self.upload_file(
-                        file_path=self.config.LOG_FILE,
-                        repo_id=self.repo_ids["log"],
-                        path_in_repo=self.config.LOG_FILE
-                    )
-                    self.upload_file(
-                        file_path=self.config.LIST_FILE,
-                        repo_id=self.repo_ids["model_list"],
-                        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_ids)+"\n")
-                        f.write(f"{self.repo_ids['current']}\n")
-                    logger.info("No new models found. Updated log.")
-                    self.upload_file(
-                        file_path=self.config.LOG_FILE,
-                        repo_id=self.repo_ids["log"],
-                        path_in_repo=self.config.LOG_FILE
-                    )
-                    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
-config = Config()
-crawler = CivitAICrawler(config)
-app = crawler.app
+        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()