diff --git a/__pycache__/main.cpython-313.pyc b/__pycache__/main.cpython-313.pyc index 12cde9d..f3c3640 100644 Binary files a/__pycache__/main.cpython-313.pyc and b/__pycache__/main.cpython-313.pyc differ diff --git a/adapters/__pycache__/Exception.cpython-313.pyc b/adapters/__pycache__/Exception.cpython-313.pyc index 6528c4b..713f69c 100644 Binary files a/adapters/__pycache__/Exception.cpython-313.pyc and b/adapters/__pycache__/Exception.cpython-313.pyc differ diff --git a/adapters/__pycache__/__init__.cpython-313.pyc b/adapters/__pycache__/__init__.cpython-313.pyc index d4646d3..60293c0 100644 Binary files a/adapters/__pycache__/__init__.cpython-313.pyc and b/adapters/__pycache__/__init__.cpython-313.pyc differ diff --git a/adapters/__pycache__/google_adapter.cpython-313.pyc b/adapters/__pycache__/google_adapter.cpython-313.pyc index c1a7312..e61c572 100644 Binary files a/adapters/__pycache__/google_adapter.cpython-313.pyc and b/adapters/__pycache__/google_adapter.cpython-313.pyc differ diff --git a/api/__pycache__/dependency.cpython-313.pyc b/api/__pycache__/dependency.cpython-313.pyc index 4b06fb7..ca98621 100644 Binary files a/api/__pycache__/dependency.cpython-313.pyc and b/api/__pycache__/dependency.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/assets_router.cpython-313.pyc b/api/endpoints/__pycache__/assets_router.cpython-313.pyc index ff27036..ab6adc2 100644 Binary files a/api/endpoints/__pycache__/assets_router.cpython-313.pyc and b/api/endpoints/__pycache__/assets_router.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/character_router.cpython-313.pyc b/api/endpoints/__pycache__/character_router.cpython-313.pyc index dbd9780..d20fbb0 100644 Binary files a/api/endpoints/__pycache__/character_router.cpython-313.pyc and b/api/endpoints/__pycache__/character_router.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/generation_router.cpython-313.pyc b/api/endpoints/__pycache__/generation_router.cpython-313.pyc index 117aef6..3c9f1b0 100644 Binary files a/api/endpoints/__pycache__/generation_router.cpython-313.pyc and b/api/endpoints/__pycache__/generation_router.cpython-313.pyc differ diff --git a/api/endpoints/assets_router.py b/api/endpoints/assets_router.py index 89985db..c996500 100644 --- a/api/endpoints/assets_router.py +++ b/api/endpoints/assets_router.py @@ -21,17 +21,31 @@ router = APIRouter(prefix="/api/assets", tags=["Assets"]) @router.get("/{asset_id}") -async def get_asset(asset_id: str, request: Request,dao: DAO = Depends(get_dao),) -> Response: - logger.debug(f"get_asset called for ID: {asset_id}") +async def get_asset( + asset_id: str, + request: Request, + thumbnail: bool = False, + dao: DAO = Depends(get_dao) +) -> Response: + logger.debug(f"get_asset called for ID: {asset_id}, thumbnail={thumbnail}") asset = await dao.assets.get_asset(asset_id) # 2. Проверка на существование if not asset: raise HTTPException(status_code=404, detail="Asset not found") + headers = { # Кэшировать на 1 год (31536000 сек) "Cache-Control": "public, max-age=31536000, immutable" } - return Response(content=asset.data, media_type="image/png", headers=headers) + + content = asset.data + media_type = "image/png" # Default, or detect + + if thumbnail and asset.thumbnail: + content = asset.thumbnail + media_type = "image/jpeg" + + return Response(content=content, media_type=media_type, headers=headers) @router.get("") @@ -41,7 +55,13 @@ async def get_assets(request: Request, dao: DAO = Depends(get_dao), limit: int = # assets = await dao.assets.get_assets() # This line seemed redundant/conflicting in original code total_count = await dao.assets.get_asset_count() - return AssetsResponse(assets=assets, total_count=total_count) + # Manually map to DTO to trigger computed fields validation if necessary, + # but primarily to ensure valid Pydantic models for the response list. + # Asset.model_dump() generally includes computed fields (url) if configured. + # Let's ensure strict conversion. + asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets] + + return AssetsResponse(assets=asset_responses, total_count=total_count) @@ -62,11 +82,16 @@ async def upload_asset( if not data: raise HTTPException(status_code=400, detail="Empty file") + # Generate thumbnail + from utils.image_utils import create_thumbnail + thumbnail_bytes = await asyncio.to_thread(create_thumbnail, data) + asset = Asset( name=file.filename or "upload", type=AssetType.IMAGE, linked_char_id=linked_char_id, data=data, + thumbnail=thumbnail_bytes ) asset_id = await dao.assets.create_asset(asset) @@ -79,4 +104,39 @@ async def upload_asset( type=asset.type.value if hasattr(asset.type, "value") else asset.type, linked_char_id=asset.linked_char_id, created_at=asset.created_at, - ) \ No newline at end of file + url=asset.url + ) + + +@router.post("/regenerate_thumbnails") +async def regenerate_thumbnails(dao: DAO = Depends(get_dao)): + """ + Regenerates thumbnails for all existing image assets that don't have one. + """ + logger.info("Starting thumbnail regeneration task") + from utils.image_utils import create_thumbnail + import asyncio + + # Get all assets (pagination loop might be needed for huge datasets, but simple list for now) + # We'll rely on DAO providing a method or just fetch large chunk. + # Assuming get_assets might have limit, let's fetch in chunks or just all if possible within limit. + # Ideally should use a specific repo method for iteration. + # For now, let's fetch first 1000 or similar. + assets = await dao.assets.get_assets(limit=1000, offset=0, with_data=True) + logger.info(f"Found {len(assets)} assets") + count = 0 + updated = 0 + + for asset in assets: + if asset.type == AssetType.IMAGE and asset.data : + try: + thumb = await asyncio.to_thread(create_thumbnail, asset.data) + if thumb: + asset.thumbnail = thumb + await dao.assets.update_asset(asset.id, asset) + updated += 1 + except Exception as e: + logger.error(f"Failed to regenerate thumbnail for asset {asset.id}: {e}") + count += 1 + + return {"status": "completed", "processed": count, "updated": updated} \ No newline at end of file diff --git a/api/endpoints/character_router.py b/api/endpoints/character_router.py index ae72b9b..2f55f5f 100644 --- a/api/endpoints/character_router.py +++ b/api/endpoints/character_router.py @@ -5,7 +5,7 @@ from pydantic import BaseModel from starlette.exceptions import HTTPException from starlette.requests import Request -from api.models.AssetDTO import AssetsResponse +from api.models.AssetDTO import AssetsResponse, AssetResponse from api.models.GenerationRequest import GenerationRequest, GenerationResponse from models.Asset import Asset from models.Character import Character @@ -35,7 +35,9 @@ async def get_character_assets(character_id: str, dao: DAO = Depends(get_dao), l raise HTTPException(status_code=404, detail="Character not found") assets = await dao.assets.get_assets_by_char_id(character_id, limit, offset) total_count = await dao.assets.get_asset_count(character_id) - return AssetsResponse(assets=assets, total_count=total_count) + + asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets] + return AssetsResponse(assets=asset_responses, total_count=total_count) @router.get("/{character_id}", response_model=Character) diff --git a/api/models/AssetDTO.py b/api/models/AssetDTO.py index 4df084e..9fefd7c 100644 --- a/api/models/AssetDTO.py +++ b/api/models/AssetDTO.py @@ -6,14 +6,14 @@ from pydantic import BaseModel from models.Asset import Asset -class AssetsResponse(BaseModel): - assets: List[Asset] - total_count: int - - class AssetResponse(BaseModel): id: str name: str type: str linked_char_id: Optional[str] = None - created_at: datetime \ No newline at end of file + created_at: datetime + url: Optional[str] = None + +class AssetsResponse(BaseModel): + assets: List[AssetResponse] + total_count: int \ No newline at end of file diff --git a/api/models/__pycache__/AssetDTO.cpython-313.pyc b/api/models/__pycache__/AssetDTO.cpython-313.pyc index 8a355fb..2bfba91 100644 Binary files a/api/models/__pycache__/AssetDTO.cpython-313.pyc and b/api/models/__pycache__/AssetDTO.cpython-313.pyc differ diff --git a/api/models/__pycache__/GenerationRequest.cpython-313.pyc b/api/models/__pycache__/GenerationRequest.cpython-313.pyc index 152561e..9bddcb9 100644 Binary files a/api/models/__pycache__/GenerationRequest.cpython-313.pyc and b/api/models/__pycache__/GenerationRequest.cpython-313.pyc differ diff --git a/api/service/__pycache__/generation_service.cpython-313.pyc b/api/service/__pycache__/generation_service.cpython-313.pyc index 62c6ea2..beca295 100644 Binary files a/api/service/__pycache__/generation_service.cpython-313.pyc and b/api/service/__pycache__/generation_service.cpython-313.pyc differ diff --git a/api/service/generation_service.py b/api/service/generation_service.py index 3bcce54..9d143da 100644 --- a/api/service/generation_service.py +++ b/api/service/generation_service.py @@ -66,7 +66,7 @@ class GenerationService: async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str: future_prompt = """You are an prompt-assistant. You improving user-entered prompts for image generation. User may upload reference image too. I will provide sources prompt entered by user. Understand user needs and generate best variation of prompt. - ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT:""" + ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT: """ future_prompt += prompt assets_data = [] if assets is not None: @@ -88,7 +88,7 @@ class GenerationService: async def get_generations(self, character_id: Optional[str] = None, limit: int = 10, offset: int = 0) -> List[ Generation]: - return await self.dao.generations.get_generations(limit=limit, offset=offset) + return await self.dao.generations.get_generations(character_id = character_id,limit=limit, offset=offset) async def get_generation(self, generation_id: str) -> Optional[GenerationResponse]: gen = await self.dao.generations.get_generation(generation_id) @@ -162,7 +162,7 @@ class GenerationService: for asset in reference_assets if asset.data is not None and asset.type == AssetType.IMAGE ) - generation_prompt+=f"PROMPT: {generation.prompt}" + generation_prompt+=f" PROMPT: {generation.prompt}" logger.info(f"Final generation prompt assembled. Length: {len(generation_prompt)}. Media count: {len(media_group_bytes)}") # 3. Запускаем процесс генерации и симуляцию прогресса @@ -209,11 +209,19 @@ class GenerationService: created_assets: List[Asset] = [] for idx, img_bytes in enumerate(generated_bytes_list): + # Generate thumbnail + thumbnail_bytes = None + # Assuming AssetType.IMAGE since we are in generated_bytes_list which are images usually + # Or use explicit check if we have distinct types in list (not currently) + from utils.image_utils import create_thumbnail + thumbnail_bytes = await asyncio.to_thread(create_thumbnail, img_bytes) + new_asset = Asset( name=f"Generated_{generation.linked_character_id}_{random.randint(1000, 9999)}", type=AssetType.IMAGE, linked_char_id=generation.linked_character_id, # Если генерация привязана к персонажу data=img_bytes, + thumbnail=thumbnail_bytes # Остальные поля заполнятся дефолтными значениями (created_at) ) @@ -236,6 +244,8 @@ class GenerationService: end_time = datetime.now() generation.execution_time_seconds = (end_time - start_time).total_seconds() + logger.info(f"DEBUG: Saving generation {generation.id}. Metrics: api_exec={generation.api_execution_time_seconds}, tokens={generation.token_usage}, exec={generation.execution_time_seconds}") + await self.dao.generations.update_generation(generation) logger.info(f"Generation {generation.id} completed successfully. {len(created_assets)} assets created. Total Time: {generation.execution_time_seconds:.2f}s") diff --git a/models/Asset.py b/models/Asset.py index 79730a2..4c27e61 100644 --- a/models/Asset.py +++ b/models/Asset.py @@ -18,6 +18,7 @@ class Asset(BaseModel): data: Optional[bytes] = None tg_doc_file_id: Optional[str] = None tg_photo_file_id: Optional[str] = None + thumbnail: Optional[bytes] = None tags: List[str] = [] created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/__pycache__/Asset.cpython-313.pyc b/models/__pycache__/Asset.cpython-313.pyc index 97a91d8..feb82a5 100644 Binary files a/models/__pycache__/Asset.cpython-313.pyc and b/models/__pycache__/Asset.cpython-313.pyc differ diff --git a/models/__pycache__/Generation.cpython-313.pyc b/models/__pycache__/Generation.cpython-313.pyc index c3f8a9f..036c27f 100644 Binary files a/models/__pycache__/Generation.cpython-313.pyc and b/models/__pycache__/Generation.cpython-313.pyc differ diff --git a/models/__pycache__/enums.cpython-313.pyc b/models/__pycache__/enums.cpython-313.pyc index 254c746..e307fb2 100644 Binary files a/models/__pycache__/enums.cpython-313.pyc and b/models/__pycache__/enums.cpython-313.pyc differ diff --git a/repos/__pycache__/assets_repo.cpython-313.pyc b/repos/__pycache__/assets_repo.cpython-313.pyc index 3da7588..317cb0e 100644 Binary files a/repos/__pycache__/assets_repo.cpython-313.pyc and b/repos/__pycache__/assets_repo.cpython-313.pyc differ diff --git a/repos/__pycache__/generation_repo.cpython-313.pyc b/repos/__pycache__/generation_repo.cpython-313.pyc index e21a10d..f64f0d8 100644 Binary files a/repos/__pycache__/generation_repo.cpython-313.pyc and b/repos/__pycache__/generation_repo.cpython-313.pyc differ diff --git a/repos/assets_repo.py b/repos/assets_repo.py index da3af33..1c00fd9 100644 --- a/repos/assets_repo.py +++ b/repos/assets_repo.py @@ -15,8 +15,12 @@ class AssetsRepo: return str(res.inserted_id) - async def get_assets(self, limit: int = 10, offset: int = 0) -> List[Asset]: - res = await self.collection.find({}, {"data": 0}).sort("created_at", -1).skip(offset).limit(limit).to_list(None) + async def get_assets(self, limit: int = 10, offset: int = 0, with_data: bool = False) -> List[Asset]: + args = {} + if not with_data: + args["data"] = 0 + args["thumbnail"] = 0 + res = await self.collection.find({}, args).sort("created_at", -1).skip(offset).limit(limit).to_list(None) assets = [] for doc in res: # Конвертируем ObjectId в строку и кладем в поле id @@ -32,6 +36,7 @@ class AssetsRepo: projection = {"_id": 1, "name": 1, "type": 1, "tg_doc_file_id": 1} if with_data: projection["data"] = 1 + projection["thumbnail"] = 1 res = await self.collection.find_one({"_id": ObjectId(asset_id)}, projection) @@ -63,7 +68,7 @@ class AssetsRepo: async def get_assets_by_ids(self, asset_ids: List[str]) -> List[Asset]: object_ids = [ObjectId(asset_id) for asset_id in asset_ids] - res = self.collection.find({"_id": {"$in": object_ids}}) + res = self.collection.find({"_id": {"$in": object_ids}}, {"thumbnail": 0}) assets = [] async for doc in res: doc["id"] = str(doc.pop("_id")) diff --git a/repos/generation_repo.py b/repos/generation_repo.py index ed5af66..24f9ad8 100644 --- a/repos/generation_repo.py +++ b/repos/generation_repo.py @@ -28,7 +28,9 @@ class GenerationRepo: limit: int = 10, offset: int = 10) -> List[Generation]: args = {} if character_id is not None: - args["character_id"] = character_id + args["linked_character_id"] = character_id + else: + args["linked_character_id"] = None if status is not None: args["status"] = status res = await self.collection.find(args).sort("created_at", -1).skip( diff --git a/routers/__pycache__/gen_router.cpython-313.pyc b/routers/__pycache__/gen_router.cpython-313.pyc index ae783b4..bb98d46 100644 Binary files a/routers/__pycache__/gen_router.cpython-313.pyc and b/routers/__pycache__/gen_router.cpython-313.pyc differ diff --git a/utils/__pycache__/image_utils.cpython-313.pyc b/utils/__pycache__/image_utils.cpython-313.pyc new file mode 100644 index 0000000..a9c80bf Binary files /dev/null and b/utils/__pycache__/image_utils.cpython-313.pyc differ diff --git a/utils/image_utils.py b/utils/image_utils.py new file mode 100644 index 0000000..480515d --- /dev/null +++ b/utils/image_utils.py @@ -0,0 +1,27 @@ +from io import BytesIO +from typing import Tuple, Optional +from PIL import Image +import logging + +logger = logging.getLogger(__name__) + +def create_thumbnail(image_data: bytes, size: Tuple[int, int] = (800, 800)) -> Optional[bytes]: + """ + Creates a thumbnail from image bytes. + Returns the thumbnail as bytes (JPEG format) or None if failed. + """ + try: + with Image.open(BytesIO(image_data)) as img: + # Convert to RGB if necessary (e.g. for RGBA/P images saving as JPEG) + if img.mode in ('RGBA', 'P'): + img = img.convert('RGB') + + img.thumbnail(size) + + thumb_io = BytesIO() + img.save(thumb_io, format='JPEG', quality=85) + thumb_io.seek(0) + return thumb_io.read() + except Exception as e: + logger.error(f"Failed to create thumbnail: {e}") + return None