diff --git a/py/services/civitai_client.py b/py/services/civitai_client.py index 60bfe463..1198716c 100644 --- a/py/services/civitai_client.py +++ b/py/services/civitai_client.py @@ -267,7 +267,7 @@ class CivitaiClient: return None, error_msg async def get_model_metadata(self, model_id: str) -> Tuple[Optional[Dict], int]: - """Fetch model metadata (description and tags) from Civitai API + """Fetch model metadata (description, tags, and creator info) from Civitai API Args: model_id: The Civitai model ID @@ -294,10 +294,14 @@ class CivitaiClient: # Extract relevant metadata metadata = { "description": data.get("description") or "No model description available", - "tags": data.get("tags", []) + "tags": data.get("tags", []), + "creator": { + "username": data.get("creator", {}).get("username"), + "image": data.get("creator", {}).get("image") + } } - if metadata["description"] or metadata["tags"]: + if metadata["description"] or metadata["tags"] or metadata["creator"]["username"]: return metadata, status_code else: logger.warning(f"No metadata found for model {model_id}") diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 47d27d5b..978211ea 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -154,7 +154,7 @@ class DownloadManager: metadata = LoraMetadata.from_civitai_info(version_info, file_info, save_path) logger.info(f"Creating LoraMetadata for {file_name}") - # 5.1 Get and update model tags and description + # 5.1 Get and update model tags, description and creator info model_id = version_info.get('modelId') if model_id: model_metadata, _ = await civitai_client.get_model_metadata(str(model_id)) @@ -163,6 +163,8 @@ class DownloadManager: metadata.tags = model_metadata.get("tags", []) if model_metadata.get("description"): metadata.modelDescription = model_metadata.get("description", "") + if model_metadata.get("creator"): + metadata.civitai["creator"] = model_metadata.get("creator") # 6. Start download process result = await self._execute_download(