Add civit model pull ability

2025-09-20 12:12:56 +07:00
parent 65bffc38eb
commit 991d655756
9 changed files with 632 additions and 13 deletions
--- a/modules/civit/client.py
+++ b/modules/civit/client.py
@@ -7,7 +7,6 @@ import time
 from typing import Optional
 from requests import Session

-from modelspace.Repository import Repository
 from pythonapp.Libs.ConfigDataClass import Config

@dataclass
@@ -108,4 +107,70 @@ class Client:
    def get_creators_raw(self, page=None, limit = 200, query = None): return self.get_creators_tags_raw('creators', page, limit, query)
    def get_tags_raw(self, page=None, limit = 200, query = None): return self.get_creators_tags_raw('tags', page, limit, query)

+    def get_model_raw(self, model_id: int):
+        try:
+            return self.make_get_request(f'{self.config.base_url}/api/v1/models/{model_id}').json()
+        except requests.exceptions.HTTPError as e:
+            print(e)
+            return {}
+
+    def download_file(self, url: str, path: str, chill_time: int = 3, max_retries: int = 3):
+        """
+        Загружает файл по URL в указанный путь
+
+        Args:
+            url (str): URL файла для загрузки
+            path (str): Путь для сохранения файла
+            chill_time (int): Время ожидания в секундах при ошибке (по умолчанию 3)
+            max_retries (int): Максимальное количество попыток загрузки (по умолчанию 3)
+        """
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        for attempt in range(max_retries):
+            try:
+                # Создаем запрос с прогресс-баром
+                response = self.session.get(url, stream=True, timeout=30)
+                response.raise_for_status()
+
+                # Получаем размер файла
+                total_size = int(response.headers.get('content-length', 0))
+
+                # Загружаем файл по частями
+                with open(path, 'wb') as file:
+                    downloaded = 0
+                    for chunk in response.iter_content(chunk_size=8192):
+                        if chunk:
+                            file.write(chunk)
+                            downloaded += len(chunk)
+
+                            # Отображаем прогресс
+                            if total_size > 0:
+                                progress = (downloaded / total_size) * 100
+                                print(f"\rDownloading: {progress:.1f}% ({downloaded}/{total_size} bytes)", end='',
+                                      flush=True)
+
+                print(f"\nFile downloaded successfully to {path}")
+                return True
+
+            except requests.exceptions.RequestException as e:
+                print(f"Attempt {attempt + 1} failed: {e}")
+
+                if attempt < max_retries - 1:
+                    print(f"Waiting {chill_time} seconds before retry...")
+                    time.sleep(chill_time)
+                else:
+                    print(f"Failed to download file after {max_retries} attempts")
+                    return False
+
+            except IOError as e:
+                print(f"IO Error while saving file: {e}")
+                return False
+
+            except Exception as e:
+                print(f"Unexpected error: {e}")
+                return False
+
+        return False
+

--- a/modules/civit/datamodel.py
+++ b/modules/civit/datamodel.py
@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from .datamodel_base import ForwardingBase
+
+@dataclass
+class Creator(ForwardingBase):
+    username: Optional[str] = None
+    modelCount: Optional[int] = None
+    link: Optional[str] = None
+    image: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {}
+        self._key_field = 'username'
+
+@dataclass
+class Tag(ForwardingBase):
+    name: Optional[str] = None
+    link: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {}
+        self._key_field = 'name'
+
+@dataclass
+class ModelVersionStats(ForwardingBase):
+    downloadCount: Optional[int] = None
+    ratingCount: Optional[int] = None
+    rating: Optional[int] = None
+    thumbsUpCount: Optional[int] = None
+    thumbsDownCount: Optional[int] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {}
+
+
+
+@dataclass
+class ModelVersion(ForwardingBase):
+    id: Optional[int] = None
+    index: Optional[int] = None
+    name: Optional[str] = None
+    baseModel: Optional[str] = None
+    baseModelType: Optional[str] = None
+    publishedAt: Optional[str] = None
+    availability: Optional[str] = None
+    nsfwLevel: Optional[int] = None
+    description: Optional[str] = None
+    trainedWords: Optional[list[str]] = None
+    stats: Optional[ModelVersionStats] = None
+    supportsGeneration: Optional[bool] = None
+    downloadUrl: Optional[str] = None
+    # FILES
+    # IMAGES
+
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {
+            'stats': ModelVersionStats,
+        }
+        self._key_field = 'id'
+
+
+
+
+@dataclass
+class ModelStats(ForwardingBase):
+    downloadCount: Optional[int] = None
+    favoriteCount: Optional[int] = None
+    thumbsUpCount: Optional[int] = None
+    thumbsDownCount: Optional[int] = None
+    commentCount: Optional[int] = None
+    ratingCount: Optional[int] = None
+    rating: Optional[int] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {}
+
+
+
+
+@dataclass
+class Model(ForwardingBase):
+    id: Optional[int] = None
+    name: Optional[str] = None
+    description: Optional[str] = None
+    allowNoCredit: Optional[bool] = None
+    allowCommercialUse: Optional[list] = None
+    allowDerivatives: Optional[bool] = None
+    allowDifferentLicense: Optional[bool] = None
+    type: Optional[str] = None
+    minor: Optional[bool] = None
+    sfwOnly: Optional[bool] = None
+    poi: Optional[bool] = None
+    nsfw: Optional[bool] = None
+    nsfwLevel: Optional[int] = None
+    availability: Optional[str] = None
+    cosmetic: Optional[str] = None
+    supportsGeneration: Optional[bool] = None
+    stats: Optional[ModelStats] = None
+    creator: Optional[Creator] = None
+    tags: Optional[list[Tag]] = None
+    modelVersions: Optional[list[ModelVersion]] = None
+
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._forwarding = {
+            'stats': ModelStats,
+            'creator': Creator,
+            'tags': Tag,
+            'modelVersions': ModelVersion,
+        }
+        self._key_field = 'id'
+
+
--- a/modules/civit/fetch.py
+++ b/modules/civit/fetch.py
@@ -352,7 +352,12 @@ class Fetch:
        print(f"Fetching {entity}...")
        path = Path(client.path) / ('fetch_' + entity)
        items = list()
-        first_page = client.make_get_request(url=f'{client.config.base_url}/api/v1/{entity}{client.build_query_string(params)}').json()
+        url = f'{client.config.base_url}/api/v1/{entity}{client.build_query_string(params)}'
+        first_page = client.make_get_request(url)
+        if not first_page:
+            with open(Path(client.path) / 'bugs.log', 'a') as f: f.write(url + '\n')
+            return items
+        first_page = first_page.json()
        if first_page.get('items', None): items.extend(first_page.get('items', None))
        if save:
            path.mkdir(exist_ok=True)
@@ -419,7 +424,7 @@ class Fetch:
                    for item in page_items: page_items_dict[item['id']] = item
                    print(f'Added {len(page_items_dict) - l} images by {sort} sort crawl. {len(page_items_dict)} images total')

-                page_items = [key for key, value in page_items_dict.items()]
+                page_items = [value for key, value in page_items_dict.items()]


            l = len(items)