freepik/freepik_cli/commands/edit.py

"""expand-image, relight, style-transfer 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.edit import EditAPI
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
from freepik_cli.utils.polling import FreepikTaskError, FreepikTimeoutError, PollConfig, poll_task

_EXPAND_MODELS = ["flux-pro", "ideogram", "seedream-v4-5"]


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[/cyan].")
        raise typer.Exit(1)
    return key


def expand_image(
    image: Annotated[
        Path,
        typer.Argument(help="Image to expand (outpaint)", exists=True),
    ],
    left: Annotated[int, typer.Option("--left", help="Pixels to add on the left (0–2048)", min=0, max=2048)] = 0,
    right: Annotated[int, typer.Option("--right", help="Pixels to add on the right (0–2048)", min=0, max=2048)] = 0,
    top: Annotated[int, typer.Option("--top", help="Pixels to add on top (0–2048)", min=0, max=2048)] = 0,
    bottom: Annotated[int, typer.Option("--bottom", help="Pixels to add on bottom (0–2048)", min=0, max=2048)] = 0,
    prompt: Annotated[Optional[str], typer.Option("--prompt", "-p", help="Optional prompt to guide the expansion")] = None,
    model: Annotated[str, typer.Option("--model", "-m", help=f"Expansion model: {', '.join(_EXPAND_MODELS)}")] = "flux-pro",
    seed: Annotated[Optional[int], typer.Option("--seed")] = None,
    output: Annotated[Optional[Path], typer.Option("--output", "-o")] = 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]Expand an image[/bold] by adding new content around its edges (outpainting).

    [dim]Examples:[/dim]
      freepik expand-image photo.jpg --left 512 --right 512 --prompt "lush forest"
      freepik expand-image banner.png --bottom 256 --model seedream-v4-5
    """
    if not any([left, right, top, bottom]):
        print_error("At least one of --left, --right, --top, --bottom must be > 0.")
        raise typer.Exit(1)

    if model not in _EXPAND_MODELS:
        print_error(f"Unknown model '{model}'.", hint=f"Choose from: {', '.join(_EXPAND_MODELS)}")
        raise typer.Exit(1)

    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)
    image_b64 = image_to_base64(image)

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

        with console.status("[info]Submitting image expansion…[/info]"):
            try:
                task_id = api.expand_submit(model, image_b64, left, right, top, bottom, prompt, seed)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "expand-image", model)
            return

        poll_config = PollConfig(task_type="expand", max_wait=config.poll_timeout)
        from freepik_cli.api.images import ImageAPI
        img_api = ImageAPI(client)
        try:
            result = poll_task(
                check_fn=lambda tid: img_api.expand_status(model, tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={"Model": model, "Expand": f"L{left} R{right} T{top} B{bottom}"},
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

        from freepik_cli.api.models import get_output_urls
        urls = get_output_urls(result)
        if not urls:
            print_error("Expansion completed but no output URL found.")
            raise typer.Exit(1)

        out = output or auto_output_path("expanded", model, 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=model, output_path=out, width=w, height=h, task_type="expand"))


def relight_image(
    image: Annotated[Path, typer.Argument(help="Image to relight", exists=True)],
    prompt: Annotated[Optional[str], typer.Option("--prompt", "-p", help="Lighting description e.g. 'golden hour sunlight'")] = None,
    style: Annotated[Optional[str], typer.Option("--style", "-s", help="Lighting style preset")] = None,
    output: Annotated[Optional[Path], typer.Option("--output", "-o")] = 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]Relight an image[/bold] using AI-controlled lighting. [dim](Premium feature)[/dim]

    [dim]Examples:[/dim]
      freepik relight portrait.jpg --prompt "dramatic studio lighting"
      freepik relight scene.png --prompt "warm golden hour" --output relit.jpg
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)
    image_b64 = image_to_base64(image)

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

        with console.status("[info]Submitting image relight…[/info]"):
            try:
                task_id = api.relight_submit(image_b64, prompt, style)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "relight", "image-relight")
            return

        poll_config = PollConfig(task_type="relight", max_wait=config.poll_timeout)
        try:
            result = poll_task(
                check_fn=lambda tid: api.relight_status(tid),
                task_id=task_id,
                config=poll_config,
                console=console,
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

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

        out = output or auto_output_path("relit", "relight", 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="image-relight", output_path=out, width=w, height=h, task_type="relight"))


def style_transfer(
    content: Annotated[Path, typer.Argument(help="Content image (what to keep)", exists=True)],
    style_image: Annotated[Path, typer.Argument(help="Style image (the artistic style to apply)", exists=True)],
    strength: Annotated[Optional[float], typer.Option("--strength", help="Style strength 0.0–1.0", min=0.0, max=1.0)] = None,
    output: Annotated[Optional[Path], typer.Option("--output", "-o")] = 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]Apply an artistic style[/bold] from one image onto another. [dim](Premium feature)[/dim]

    [dim]Examples:[/dim]
      freepik style-transfer photo.jpg painting.jpg --strength 0.8
      freepik style-transfer portrait.png van_gogh.jpg --output styled.jpg
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)

    content_b64 = image_to_base64(content)
    style_b64 = image_to_base64(style_image)

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

        with console.status("[info]Submitting style transfer…[/info]"):
            try:
                task_id = api.style_transfer_submit(content_b64, style_b64, strength)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "style-transfer", "image-style-transfer")
            return

        poll_config = PollConfig(task_type="style-transfer", max_wait=config.poll_timeout)
        try:
            result = poll_task(
                check_fn=lambda tid: api.style_transfer_status(tid),
                task_id=task_id,
                config=poll_config,
                console=console,
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

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

        out = output or auto_output_path("styled", "style-transfer", "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="image-style-transfer", output_path=out, width=w, height=h, task_type="style-transfer"))