forge/backend/app/services/video_upscaler.py

"""Video Upscaler Service - Topaz Labs API"""
import httpx
import os
from uuid import uuid4
from datetime import datetime
import asyncio

from app.database import SessionLocal
from app.models.job import Job
from app.models.asset import Asset
from app.config import settings


async def upscale(job_id: str):
    """Upscale video using Topaz Labs API"""
    db = SessionLocal()
    try:
        job = db.query(Job).filter(Job.id == job_id).first()
        if not job:
            return

        input_data = job.input_data
        input_asset_ids = job.input_asset_ids

        if not input_asset_ids:
            raise ValueError("No input asset provided")

        input_asset = db.query(Asset).filter(Asset.id == input_asset_ids[0]).first()
        if not input_asset:
            raise ValueError("Input asset not found")

        job.progress = 5
        job.api_provider = "topaz"
        job.api_model = input_data.get("model", "auto")
        db.commit()

        scale = input_data.get("scale", 2)
        model = input_data.get("model", "auto")
        frame_interpolation = input_data.get("frame_interpolation", 1)

        # Get video info (simplified - would need ffprobe in production)
        video_info = {
            "container": "mp4",
            "size": input_asset.file_size_bytes,
            "duration": float(input_asset.duration_seconds or 10),
            "frameCount": int((input_asset.duration_seconds or 10) * 30),
            "frameRate": 30,
            "resolution": {
                "width": input_asset.width or 1920,
                "height": input_asset.height or 1080
            }
        }

        output_width = video_info["resolution"]["width"] * scale
        output_height = video_info["resolution"]["height"] * scale

        job.progress = 10
        db.commit()

        async with httpx.AsyncClient(timeout=1800) as client:
            # Create video enhancement request
            response = await client.post(
                "https://api.topazlabs.com/video/v1/enhance",
                headers={
                    "X-API-Key": settings.topaz_api_key,
                    "Content-Type": "application/json"
                },
                json={
                    "source": video_info,
                    "filters": [
                        {
                            "model": model if model != "auto" else "prob-4",
                            "videoType": "Progressive",
                            "auto": "Auto" if model == "auto" else None
                        }
                    ],
                    "output": {
                        "resolution": {
                            "width": output_width,
                            "height": output_height
                        },
                        "frameRate": video_info["frameRate"] * frame_interpolation,
                        "audioCodec": "AAC",
                        "audioTransfer": "Copy",
                        "container": "mp4"
                    }
                }
            )
            response.raise_for_status()
            result = response.json()

            request_id = result.get("requestId")

            job.progress = 15
            job.api_request_id = request_id
            db.commit()

            # Accept the request and get upload URLs
            accept_response = await client.patch(
                f"https://api.topazlabs.com/video/v1/enhance/{request_id}/accept",
                headers={"X-API-Key": settings.topaz_api_key}
            )
            accept_data = accept_response.json()
            upload_urls = accept_data.get("urls", [])

            job.progress = 20
            db.commit()

            # Upload video file in parts
            with open(input_asset.file_path, "rb") as f:
                video_data = f.read()

            part_size = len(video_data) // len(upload_urls) if upload_urls else len(video_data)
            upload_results = []

            for i, url in enumerate(upload_urls):
                start = i * part_size
                end = start + part_size if i < len(upload_urls) - 1 else len(video_data)
                part_data = video_data[start:end]

                upload_response = await client.put(
                    url,
                    content=part_data,
                    headers={"Content-Type": "application/octet-stream"}
                )

                etag = upload_response.headers.get("ETag", "").strip('"')
                upload_results.append({
                    "partNum": i + 1,
                    "eTag": etag
                })

                job.progress = 20 + (i + 1) * (30 / len(upload_urls))
                db.commit()

            # Complete the upload
            await client.patch(
                f"https://api.topazlabs.com/video/v1/enhance/{request_id}/complete-upload/",
                headers={
                    "X-API-Key": settings.topaz_api_key,
                    "Content-Type": "application/json"
                },
                json={"uploadResults": upload_results}
            )

            job.progress = 50
            db.commit()

            # Poll for completion
            for _ in range(360):  # Wait up to 12 minutes
                await asyncio.sleep(2)

                status_response = await client.get(
                    f"https://api.topazlabs.com/video/v1/enhance/{request_id}/status",
                    headers={"X-API-Key": settings.topaz_api_key}
                )
                status_data = status_response.json()
                status = status_data.get("status", "")

                if status == "completed":
                    output_url = status_data.get("outputUrl")
                    if output_url:
                        video_response = await client.get(output_url)
                        upscaled_data = video_response.content

                        # Save output
                        filename = f"upscaled_{uuid4()}.mp4"
                        storage_path = os.path.join(settings.storage_path, "videos")
                        os.makedirs(storage_path, exist_ok=True)
                        file_path = os.path.join(storage_path, filename)

                        with open(file_path, "wb") as f:
                            f.write(upscaled_data)

                        # Create output asset
                        output_asset = Asset(
                            user_id=job.user_id,
                            project_id=job.project_id,
                            original_filename=filename,
                            stored_filename=filename,
                            file_path=file_path,
                            file_type="video",
                            mime_type="video/mp4",
                            file_size_bytes=len(upscaled_data),
                            width=output_width,
                            height=output_height,
                            duration_seconds=input_asset.duration_seconds,
                            source_module="video_upscaler",
                            source_job_id=job.id,
                            parent_asset_id=input_asset.id,
                            metadata={
                                "scale": scale,
                                "model": model,
                                "frame_interpolation": frame_interpolation
                            }
                        )
                        db.add(output_asset)
                        db.commit()
                        db.refresh(output_asset)

                        job.output_asset_ids = [output_asset.id]
                        job.output_data = {"asset_id": str(output_asset.id), "file_path": file_path}
                        break

                elif status == "failed":
                    raise ValueError(f"Video enhancement failed: {status_data.get('error')}")

                job.progress = min(50 + (_ * 0.14), 95)
                db.commit()

        job.progress = 100
        job.status = "completed"
        job.completed_at = datetime.utcnow()
        db.commit()

    except Exception as e:
        job.status = "failed"
        job.error_message = str(e)
        db.commit()
    finally:
        db.close()