freepik/freepik_cli/commands/generate.py

"""generate-image, generate-video, generate-icon 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.api.images import ImageAPI
from freepik_cli.api.models import IconStyle, ImageModel, VideoModel
from freepik_cli.api.videos import VideoAPI
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


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 the [cyan]FREEPIK_API_KEY[/cyan] environment variable, "
                 "or pass [cyan]--api-key YOUR_KEY[/cyan].",
        )
        raise typer.Exit(1)
    return key


def generate_image(
    prompt: Annotated[str, typer.Argument(help="Text prompt describing the image to generate")],
    model: Annotated[
        ImageModel,
        typer.Option("--model", "-m", help="AI model to use for generation", show_default=True),
    ] = ImageModel.FLUX_2_PRO,
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Output file path (auto-generated if omitted)"),
    ] = None,
    aspect_ratio: Annotated[
        Optional[str],
        typer.Option(
            "--aspect-ratio", "-a",
            help="Aspect ratio e.g. [cyan]16:9[/cyan], [cyan]1:1[/cyan], [cyan]9:16[/cyan], [cyan]4:3[/cyan]",
        ),
    ] = None,
    negative_prompt: Annotated[
        Optional[str],
        typer.Option("--negative-prompt", "-n", help="Concepts to exclude from the image"),
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Random seed for reproducibility"),
    ] = None,
    input_image: Annotated[
        Optional[Path],
        typer.Option(
            "--input-image", "-i",
            help="Reference image for img2img / editing (flux-kontext-pro)",
            exists=True, file_okay=True, dir_okay=False,
        ),
    ] = 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]Generate an image[/bold] using Freepik AI models.

    [dim]Examples:[/dim]
      freepik generate-image "a cat on the moon" --model flux-2-pro
      freepik generate-image "cyberpunk city at night" --model mystic --aspect-ratio 16:9
      freepik generate-image "make the sky orange" --model flux-kontext-pro --input-image photo.jpg
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)

    # Build request payload
    payload: dict = {"prompt": prompt}
    if aspect_ratio:
        payload["aspect_ratio"] = aspect_ratio
    if negative_prompt:
        payload["negative_prompt"] = negative_prompt
    if seed is not None:
        payload["seed"] = seed
    if input_image:
        payload["image"] = image_to_base64(input_image)

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

        with console.status(f"[info]Submitting {model.value} generation…[/info]"):
            try:
                task_id = api.generate(model, payload)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "image", model.value)
            return

        poll_config = PollConfig(task_type="image", max_wait=config.poll_timeout)
        try:
            result = poll_task(
                check_fn=lambda tid: api.get_status(model, tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={"Model": f"[magenta]{model.value}[/magenta]"},
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

        urls = api.get_output_urls(result)
        if not urls:
            print_error("Generation completed but no output URLs found.", hint="Check the Freepik dashboard.")
            raise typer.Exit(1)

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


def generate_video(
    image: Annotated[
        Path,
        typer.Argument(
            help="Source image path to animate",
            exists=True, file_okay=True, dir_okay=False,
        ),
    ],
    model: Annotated[
        VideoModel,
        typer.Option("--model", "-m", help="Video AI model", show_default=True),
    ] = VideoModel.KLING_O1_PRO,
    prompt: Annotated[
        Optional[str],
        typer.Option("--prompt", "-p", help="Motion/style guidance prompt"),
    ] = None,
    duration: Annotated[
        int,
        typer.Option("--duration", "-d", help="Video duration in seconds: [cyan]5[/cyan] or [cyan]10[/cyan]", min=5, max=10),
    ] = 5,
    aspect_ratio: Annotated[
        str,
        typer.Option("--aspect-ratio", "-a", help="Output aspect ratio: [cyan]16:9[/cyan] | [cyan]9:16[/cyan] | [cyan]1:1[/cyan]"),
    ] = "16:9",
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Random seed for reproducibility"),
    ] = None,
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Output video 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]Generate a video[/bold] from a source image using AI.

    [dim]Examples:[/dim]
      freepik generate-video photo.jpg --prompt "gentle ocean waves" --model kling-o1-pro
      freepik generate-video portrait.png --model minimax-hailuo --duration 10 --aspect-ratio 9:16
    """
    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 = VideoAPI(client)

        with console.status(f"[info]Submitting {model.value} video generation…[/info]"):
            try:
                task_id = api.generate(
                    model=model,
                    image_b64=image_b64,
                    prompt=prompt,
                    duration=duration,
                    aspect_ratio=aspect_ratio,
                    seed=seed,
                )
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "video", model.value)
            return

        poll_config = PollConfig(
            task_type="video",
            initial_delay=5.0,
            max_wait=config.poll_timeout,
        )
        try:
            result = poll_task(
                check_fn=lambda tid: api.get_status(model, tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={
                    "Model": f"[magenta]{model.value}[/magenta]",
                    "Duration": f"{duration}s",
                    "Ratio": aspect_ratio,
                },
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

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

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

        print_result(
            GenerationResult(
                task_id=task_id,
                model=model.value,
                output_path=out,
                duration=str(duration),
                task_type="video",
            )
        )


def generate_icon(
    prompt: Annotated[str, typer.Argument(help="Text prompt for the icon")],
    style: Annotated[
        IconStyle,
        typer.Option("--style", "-s", help="Icon style", show_default=True),
    ] = IconStyle.COLOR,
    steps: Annotated[
        int,
        typer.Option("--steps", help="Inference steps (10–50)", min=10, max=50),
    ] = 30,
    guidance: Annotated[
        float,
        typer.Option("--guidance", help="Guidance scale (0–10)", min=0.0, max=10.0),
    ] = 7.5,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Random seed"),
    ] = None,
    fmt: Annotated[
        str,
        typer.Option("--format", "-f", help="Output format: [cyan]png[/cyan] | [cyan]svg[/cyan]"),
    ] = "png",
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Output 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]Generate an icon[/bold] from a text prompt.

    [dim]Examples:[/dim]
      freepik generate-icon "shopping cart" --style solid --format svg
      freepik generate-icon "rocket ship" --style color --format png
    """
    config = FreepikConfig.load()
    key = _get_api_key(api_key, config)

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

        with console.status("[info]Submitting icon generation…[/info]"):
            try:
                task_id = api.generate_icon(prompt, style, steps, guidance, seed)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not wait:
            print_no_wait(task_id, "icon", f"text-to-icon/{style.value}")
            return

        poll_config = PollConfig(task_type="icon", max_wait=config.poll_timeout)
        try:
            poll_task(
                check_fn=lambda tid: api.icon_status(tid),
                task_id=task_id,
                config=poll_config,
                console=console,
                extra_info={"Style": style.value, "Format": fmt},
            )
        except (FreepikTaskError, FreepikTimeoutError) as exc:
            print_error(str(exc))
            raise typer.Exit(1)

        # Render to get download URL
        with console.status(f"[info]Rendering icon as {fmt.upper()}…[/info]"):
            try:
                url = api.render_icon(task_id, fmt)
            except FreepikAPIError as exc:
                print_error(str(exc))
                raise typer.Exit(1)

        if not url:
            print_error("Icon generated but render URL not found.")
            raise typer.Exit(1)

        out = output or auto_output_path("icon", style.value, fmt, config.default_output_dir)
        save_from_url(url, out, console)

        print_result(
            GenerationResult(
                task_id=task_id,
                model=f"text-to-icon/{style.value}",
                output_path=out,
                task_type="icon",
            )
        )