feat: Implement video generation functionality and integrate with Kling API.
This commit is contained in:
@@ -3,6 +3,7 @@ from fastapi import Request, Depends
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
|
||||
from adapters.google_adapter import GoogleAdapter
|
||||
from adapters.kling_adapter import KlingAdapter
|
||||
from api.service.generation_service import GenerationService
|
||||
from repos.dao import DAO
|
||||
|
||||
@@ -36,14 +37,18 @@ def get_dao(
|
||||
# так что DAO создастся один раз за запрос.
|
||||
return DAO(mongo_client, s3_adapter)
|
||||
|
||||
def get_kling_adapter(request: Request) -> Optional[KlingAdapter]:
|
||||
return request.app.state.kling_adapter
|
||||
|
||||
# Провайдер сервиса (собирается из DAO и Gemini)
|
||||
def get_generation_service(
|
||||
dao: DAO = Depends(get_dao),
|
||||
gemini: GoogleAdapter = Depends(get_gemini_client),
|
||||
s3_adapter: S3Adapter = Depends(get_s3_adapter),
|
||||
bot: Bot = Depends(get_bot_client),
|
||||
kling_adapter: Optional[KlingAdapter] = Depends(get_kling_adapter),
|
||||
) -> GenerationService:
|
||||
return GenerationService(dao, gemini, s3_adapter, bot)
|
||||
return GenerationService(dao, gemini, s3_adapter, bot, kling_adapter=kling_adapter)
|
||||
|
||||
from fastapi import Header
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ from api.dependency import get_generation_service, get_project_id, get_dao
|
||||
from repos.dao import DAO
|
||||
|
||||
from api.models.GenerationRequest import GenerationResponse, GenerationRequest, GenerationsResponse, PromptResponse, PromptRequest
|
||||
from api.models.VideoGenerationRequest import VideoGenerationRequest
|
||||
from api.service.generation_service import GenerationService
|
||||
from models.Generation import Generation
|
||||
|
||||
@@ -113,6 +114,25 @@ async def get_running_generations(request: Request,
|
||||
return await generation_service.get_running_generations(user_id=user_id_filter, project_id=project_id)
|
||||
|
||||
|
||||
@router.post("/video/_run", response_model=GenerationResponse)
|
||||
async def post_video_generation(
|
||||
video_request: VideoGenerationRequest,
|
||||
request: Request,
|
||||
generation_service: GenerationService = Depends(get_generation_service),
|
||||
current_user: dict = Depends(get_current_user),
|
||||
project_id: Optional[str] = Depends(get_project_id),
|
||||
dao: DAO = Depends(get_dao),
|
||||
) -> GenerationResponse:
|
||||
"""Start image-to-video generation using Kling AI."""
|
||||
logger.info(f"post_video_generation called. AssetId: {video_request.image_asset_id}, Duration: {video_request.duration}s, Mode: {video_request.mode}")
|
||||
|
||||
if project_id:
|
||||
project = await dao.projects.get_project(project_id)
|
||||
if not project or str(current_user["_id"]) not in project.members:
|
||||
raise HTTPException(status_code=403, detail="Project access denied")
|
||||
video_request.project_id = project_id
|
||||
|
||||
return await generation_service.create_video_generation_task(video_request, user_id=str(current_user.get("_id")))
|
||||
|
||||
|
||||
@router.post("/import", response_model=GenerationResponse)
|
||||
|
||||
@@ -27,6 +27,7 @@ class GenerationsResponse(BaseModel):
|
||||
class GenerationResponse(BaseModel):
|
||||
id: str
|
||||
status: GenerationStatus
|
||||
gen_type: GenType = GenType.IMAGE
|
||||
failed_reason: Optional[str] = None
|
||||
|
||||
linked_character_id: Optional[str] = None
|
||||
@@ -45,6 +46,10 @@ class GenerationResponse(BaseModel):
|
||||
progress: int = 0
|
||||
cost: Optional[float] = None
|
||||
created_by: Optional[str] = None
|
||||
# Video-specific
|
||||
kling_task_id: Optional[str] = None
|
||||
video_duration: Optional[int] = None
|
||||
video_mode: Optional[str] = None
|
||||
created_at: datetime = datetime.now(UTC)
|
||||
updated_at: datetime = datetime.now(UTC)
|
||||
|
||||
|
||||
16
api/models/VideoGenerationRequest.py
Normal file
16
api/models/VideoGenerationRequest.py
Normal file
@@ -0,0 +1,16 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class VideoGenerationRequest(BaseModel):
|
||||
prompt: str = ""
|
||||
negative_prompt: Optional[str] = ""
|
||||
image_asset_id: str # ID ассета-картинки для source image
|
||||
duration: int = 5 # 5 or 10 seconds
|
||||
mode: str = "std" # "std" or "pro"
|
||||
model_name: str = "kling-v2-1"
|
||||
cfg_scale: float = 0.5
|
||||
aspect_ratio: str = "16:9"
|
||||
linked_character_id: Optional[str] = None
|
||||
project_id: Optional[str] = None
|
||||
@@ -11,7 +11,9 @@ from aiogram import Bot
|
||||
from aiogram.types import BufferedInputFile
|
||||
from adapters.Exception import GoogleGenerationException
|
||||
from adapters.google_adapter import GoogleAdapter
|
||||
from adapters.kling_adapter import KlingAdapter, KlingApiException
|
||||
from api.models.GenerationRequest import GenerationRequest, GenerationResponse, GenerationsResponse
|
||||
from api.models.VideoGenerationRequest import VideoGenerationRequest
|
||||
# Импортируйте ваши модели DAO, Asset, Generation корректно
|
||||
from models.Asset import Asset, AssetType, AssetContentType
|
||||
from models.Generation import Generation, GenerationStatus
|
||||
@@ -64,11 +66,12 @@ async def generate_image_task(
|
||||
return images_bytes, metrics
|
||||
|
||||
class GenerationService:
|
||||
def __init__(self, dao: DAO, gemini: GoogleAdapter, s3_adapter: S3Adapter, bot: Optional[Bot] = None):
|
||||
def __init__(self, dao: DAO, gemini: GoogleAdapter, s3_adapter: S3Adapter, bot: Optional[Bot] = None, kling_adapter: Optional[KlingAdapter] = None):
|
||||
self.dao = dao
|
||||
self.gemini = gemini
|
||||
self.s3_adapter = s3_adapter
|
||||
self.bot = bot
|
||||
self.kling_adapter = kling_adapter
|
||||
|
||||
|
||||
async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str:
|
||||
@@ -428,6 +431,168 @@ class GenerationService:
|
||||
|
||||
return generation
|
||||
|
||||
# === VIDEO GENERATION (Kling) ===
|
||||
|
||||
async def create_video_generation_task(self, request: VideoGenerationRequest, user_id: Optional[str] = None) -> GenerationResponse:
|
||||
"""Create a video generation task (async, returns immediately)."""
|
||||
if not self.kling_adapter:
|
||||
raise Exception("Kling adapter is not configured")
|
||||
|
||||
generation = Generation(
|
||||
status=GenerationStatus.RUNNING,
|
||||
gen_type=GenType.VIDEO,
|
||||
linked_character_id=request.linked_character_id,
|
||||
aspect_ratio=AspectRatios.SIXTEENNINE, # default for video
|
||||
quality=Quality.ONEK,
|
||||
prompt=request.prompt,
|
||||
assets_list=[request.image_asset_id],
|
||||
video_duration=request.duration,
|
||||
video_mode=request.mode,
|
||||
project_id=request.project_id,
|
||||
)
|
||||
if user_id:
|
||||
generation.created_by = user_id
|
||||
|
||||
gen_id = await self.dao.generations.create_generation(generation)
|
||||
generation.id = gen_id
|
||||
|
||||
async def runner(gen, req):
|
||||
logger.info(f"Starting background video generation task for ID: {gen.id}")
|
||||
try:
|
||||
await self.create_video_generation(gen, req)
|
||||
logger.info(f"Background video generation task finished for ID: {gen.id}")
|
||||
except Exception:
|
||||
try:
|
||||
db_gen = await self.dao.generations.get_generation(gen.id)
|
||||
if db_gen and db_gen.status != GenerationStatus.FAILED:
|
||||
db_gen.status = GenerationStatus.FAILED
|
||||
db_gen.updated_at = datetime.now(UTC)
|
||||
await self.dao.generations.update_generation(db_gen)
|
||||
except Exception:
|
||||
logger.exception("Failed to mark video generation as FAILED")
|
||||
logger.exception("create_video_generation task failed")
|
||||
|
||||
asyncio.create_task(runner(generation, request))
|
||||
return GenerationResponse(**generation.model_dump())
|
||||
|
||||
async def create_video_generation(self, generation: Generation, request: VideoGenerationRequest):
|
||||
"""Background video generation: call Kling API, poll, download result, save asset."""
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
# 1. Get source image presigned URL
|
||||
asset = await self.dao.assets.get_asset(request.image_asset_id)
|
||||
if not asset:
|
||||
raise Exception(f"Asset {request.image_asset_id} not found")
|
||||
|
||||
if not asset.minio_object_name:
|
||||
raise Exception(f"Asset {request.image_asset_id} has no S3 object")
|
||||
|
||||
presigned_url = await self.s3_adapter.get_presigned_url(asset.minio_object_name, expiration=3600)
|
||||
if not presigned_url:
|
||||
raise Exception("Failed to generate presigned URL for source image")
|
||||
|
||||
logger.info(f"Video gen {generation.id}: got presigned URL for asset {request.image_asset_id}")
|
||||
|
||||
# 2. Create Kling task
|
||||
task_data = await self.kling_adapter.create_video_task(
|
||||
image_url=presigned_url,
|
||||
prompt=request.prompt,
|
||||
negative_prompt=request.negative_prompt or "",
|
||||
model_name=request.model_name,
|
||||
duration=request.duration,
|
||||
mode=request.mode,
|
||||
cfg_scale=request.cfg_scale,
|
||||
aspect_ratio=request.aspect_ratio,
|
||||
)
|
||||
|
||||
task_id = task_data.get("task_id")
|
||||
generation.kling_task_id = task_id
|
||||
await self.dao.generations.update_generation(generation)
|
||||
|
||||
logger.info(f"Video gen {generation.id}: Kling task created, task_id={task_id}")
|
||||
|
||||
# 3. Poll for completion with progress updates
|
||||
async def progress_callback(progress_pct: int):
|
||||
generation.progress = progress_pct
|
||||
generation.updated_at = datetime.now(UTC)
|
||||
await self.dao.generations.update_generation(generation)
|
||||
|
||||
result = await self.kling_adapter.wait_for_completion(
|
||||
task_id=task_id,
|
||||
poll_interval=10,
|
||||
timeout=600,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
# 4. Extract video URL and download
|
||||
works = result.get("task_result", {}).get("videos", [])
|
||||
if not works:
|
||||
raise Exception("No video in Kling result")
|
||||
|
||||
video_url = works[0].get("url")
|
||||
video_duration = works[0].get("duration", request.duration)
|
||||
if not video_url:
|
||||
raise Exception("No video URL in Kling result")
|
||||
|
||||
logger.info(f"Video gen {generation.id}: downloading video from {video_url}")
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
video_response = await client.get(video_url)
|
||||
video_response.raise_for_status()
|
||||
video_bytes = video_response.content
|
||||
|
||||
logger.info(f"Video gen {generation.id}: downloaded {len(video_bytes)} bytes")
|
||||
|
||||
# 5. Upload to S3
|
||||
filename = f"generated_video/{generation.linked_character_id or 'no_char'}/{datetime.now().strftime('%Y%m%d_%H%M%S')}_{random.randint(1000, 9999)}.mp4"
|
||||
await self.s3_adapter.upload_file(filename, video_bytes, content_type="video/mp4")
|
||||
|
||||
# 6. Create Asset
|
||||
new_asset = Asset(
|
||||
name=f"Video_{generation.linked_character_id or 'gen'}",
|
||||
type=AssetType.GENERATED,
|
||||
content_type=AssetContentType.VIDEO,
|
||||
linked_char_id=generation.linked_character_id,
|
||||
data=None,
|
||||
minio_object_name=filename,
|
||||
minio_bucket=self.s3_adapter.bucket_name,
|
||||
thumbnail=None, # видео thumbnails можно добавить позже
|
||||
created_by=generation.created_by,
|
||||
project_id=generation.project_id,
|
||||
)
|
||||
|
||||
asset_id = await self.dao.assets.create_asset(new_asset)
|
||||
new_asset.id = str(asset_id)
|
||||
|
||||
# 7. Finalize generation
|
||||
end_time = datetime.now()
|
||||
generation.result_list = [new_asset.id]
|
||||
generation.result = new_asset.id
|
||||
generation.status = GenerationStatus.DONE
|
||||
generation.progress = 100
|
||||
generation.video_duration = video_duration
|
||||
generation.execution_time_seconds = (end_time - start_time).total_seconds()
|
||||
generation.updated_at = datetime.now(UTC)
|
||||
await self.dao.generations.update_generation(generation)
|
||||
|
||||
logger.info(f"Video generation {generation.id} completed. Asset: {new_asset.id}, Time: {generation.execution_time_seconds:.1f}s")
|
||||
|
||||
except KlingApiException as e:
|
||||
logger.error(f"Kling API error for generation {generation.id}: {e}")
|
||||
generation.status = GenerationStatus.FAILED
|
||||
generation.failed_reason = str(e)
|
||||
generation.updated_at = datetime.now(UTC)
|
||||
await self.dao.generations.update_generation(generation)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Video generation {generation.id} failed: {e}")
|
||||
generation.status = GenerationStatus.FAILED
|
||||
generation.failed_reason = str(e)
|
||||
generation.updated_at = datetime.now(UTC)
|
||||
await self.dao.generations.update_generation(generation)
|
||||
raise
|
||||
|
||||
async def delete_generation(self, generation_id: str) -> bool:
|
||||
"""
|
||||
Soft delete generation by marking it as deleted.
|
||||
|
||||
Reference in New Issue
Block a user