myapp/cli/main.py

import click
import dotenv
from huggingface_hub import InferenceClient

from myapp.palette import extract_color_clusters, generate_palette_image
from myapp.params import ImageParamType

dotenv.load_dotenv()
client = InferenceClient()


@click.group()
def cli():
    pass


@cli.command()
@click.option("--prompt", required=True)
@click.option("--target", type=click.Path(dir_okay=False), required=True)
@click.option("--model", default="black-forest-labs/FLUX.1-schnell")
@click.option("--width", default=400)
@click.option("--height", default=400)
def generate_image(prompt, target, model, width, height):
    image = client.text_to_image(
        prompt=prompt,
        model=model,
        width=width,
        height=height,
    )

    image.save(target)


@cli.command()
@click.option("--image", type=ImageParamType(), required=True)
@click.option("--target", type=click.Path(dir_okay=False), required=True)
@click.option("--n-colors", default=4)
@click.option("--shade", "shades", default=(0.0,), multiple=True)
def generate_palette(image, target, n_colors, shades):
    k_means = extract_color_clusters(image, n_colors)
    palette = generate_palette_image(k_means, shades=shades)
    palette.save(target)