freepik/freepik_cli/commands/upscale.py

"""upscale-image and upscale-video commands."""

from __future__ import annotations

from pathlib import Path
from typing import Annotated, Optional

import typer

from freepik_cli.api.client import FreepikAPIError, FreepikClient
from freepik_cli.api.models import UpscaleMode, VideoUpscaleMode
from freepik_cli.api.upscale import UpscaleAPI
from freepik_cli.utils.config import FreepikConfig
from freepik_cli.utils.console import GenerationResult, console, print_error, print_no_wait, print_result
from freepik_cli.utils.files import auto_output_path, get_image_dimensions, image_to_base64, save_from_url, video_to_base64
from freepik_cli.utils.polling import FreepikTaskError, FreepikTimeoutError, PollConfig, poll_task


def _get_api_key(api_key: Optional[str], config: FreepikConfig) -> str:
    key = api_key or config.api_key
    if not key:
        print_error(
            "No API key found.",
            hint="Set [cyan]FREEPIK_API_KEY[/cyan] or pass [cyan]--api-key YOUR_KEY[/cyan].",
        )
        raise typer.Exit(1)
    return key


def upscale_image(
    image: Annotated[
        Path,
        typer.Argument(
            help="Image file to upscale",
            exists=True, file_okay=True, dir_okay=False,
        ),
    ],
    mode: Annotated[
        UpscaleMode,
        typer.Option(
            "--mode",
            help="Upscaling mode:\n"
                 "  [cyan]precision-v2[/cyan] — faithful, detail-preserving (recommended)\n"
                 "  [cyan]precision[/cyan]   — faithful upscaling\n"
                 "  [cyan]creative[/cyan]    — AI-enhanced creative reinterpretation",
            show_default=True,
        ),
    ] = UpscaleMode.PRECISION_V2,
    scale: Annotated[
        str,
        typer.Option("--scale", help="Scale factor: [cyan]2x[/cyan] or [cyan]4x[/cyan]"),
    ] = "2x",
    creativity: Annotated[
        Optional[int],
        typer.Option(
            "--creativity",
            help="Creative enhancement level 0–10 ([cyan]creative[/cyan] mode only)",
            min=0, max=10,
        ),
    ] = None,
    prompt: Annotated[
        Optional[str],
        typer.Option("--prompt", "-p", help="Enhancement guidance ([cyan]creative[/cyan] mode only)"),
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Random seed for reproducibility"),
    ] = None,
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Output file path"),
    ] = None,
    wait: Annotated[
        bool,
        typer.Option("--wait/--no-wait", help="Wait for completion or return task ID immediately"),
    ] = True,
    api_key: Annotated[
        Optional[str],
        typer.Option("--api-key", envvar="FREEPIK_API_KEY", help="Freepik API key"),
    ] = None,
) -> None:
    """
    [bold]Upscale and enhance an image[/bold] using AI.

    [dim]Examples:[/dim]
      freepik upscale-image photo.jpg --mode precision-v2 --scale 4x
      freepik upscale-image portrait.png --mode creative --creativity 7 --prompt "sharp cinematic"
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)

    if mode != UpscaleMode.CREATIVE and (creativity is not None or prompt):
        print_error(
            "--creativity and --prompt are only supported with --mode creative.",
            hint="Switch to [cyan]--mode creative[/cyan] or remove those options.",
        )
        raise typer.Exit(1)

    image_b64 = image_to_base64(image)

    with FreepikClient(key, base_url=config.base_url) as client:
        api = UpscaleAPI(client)

        with console.status(f"[info]Submitting {mode.value} upscale ({scale})…[/info]"):
            try:
                task_id = api.upscale_image(
                    mode=mode,
                    image_b64=image_b64,
                    scale_factor=scale,
                    creativity=creativity,
                    prompt=prompt,
                    seed=seed,
                )
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "upscale-image", mode.value)
            return

        poll_config = PollConfig(task_type="upscale-image", max_wait=config.poll_timeout)
        try:
            result = poll_task(
                check_fn=lambda tid: api.upscale_image_status(mode, tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={"Mode": mode.value, "Scale": scale},
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

        urls = api.get_output_urls(result)
        if not urls:
            print_error("Upscaling completed but no output URL found.")
            raise typer.Exit(1)

        out = output or auto_output_path("upscaled", mode.value, image.suffix.lstrip(".") or "jpg", config.default_output_dir)
        save_from_url(urls[0], out, console)

        w, h = get_image_dimensions(out)
        print_result(
            GenerationResult(
                task_id=task_id,
                model=mode.value,
                output_path=out,
                width=w,
                height=h,
                task_type="upscale-image",
            )
        )


def upscale_video(
    video: Annotated[
        Path,
        typer.Argument(
            help="Video file to upscale",
            exists=True, file_okay=True, dir_okay=False,
        ),
    ],
    mode: Annotated[
        VideoUpscaleMode,
        typer.Option(
            "--mode",
            help="Upscaling mode: [cyan]standard[/cyan] | [cyan]turbo[/cyan] (faster)",
            show_default=True,
        ),
    ] = VideoUpscaleMode.STANDARD,
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Output video file path"),
    ] = None,
    wait: Annotated[
        bool,
        typer.Option("--wait/--no-wait"),
    ] = True,
    api_key: Annotated[
        Optional[str],
        typer.Option("--api-key", envvar="FREEPIK_API_KEY"),
    ] = None,
) -> None:
    """
    [bold]Upscale a video[/bold] to higher resolution using AI.

    [dim]Examples:[/dim]
      freepik upscale-video clip.mp4 --mode standard
      freepik upscale-video clip.mp4 --mode turbo --output clip_4k.mp4
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)

    video_b64 = video_to_base64(video)

    with FreepikClient(key, base_url=config.base_url) as client:
        api = UpscaleAPI(client)

        with console.status(f"[info]Submitting video upscale ({mode.value})…[/info]"):
            try:
                task_id = api.upscale_video(mode=mode, video_b64=video_b64)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "upscale-video", mode.value)
            return

        poll_config = PollConfig(
            task_type="upscale-video",
            initial_delay=5.0,
            max_wait=config.poll_timeout,
        )
        try:
            result = poll_task(
                check_fn=lambda tid: api.upscale_video_status(tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={"Mode": mode.value},
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

        urls = api.get_output_urls(result)
        if not urls:
            print_error("Upscaling completed but no output URL found.")
            raise typer.Exit(1)

        out = output or auto_output_path("upscaled_video", mode.value, "mp4", config.default_output_dir)
        save_from_url(urls[0], out, console)

        print_result(
            GenerationResult(
                task_id=task_id,
                model=mode.value,
                output_path=out,
                task_type="upscale-video",
            )
        )