myapp/cli.py

import click
import dotenv
import segno
from huggingface_hub import InferenceClient
from qrcode_artistic import write_artistic
from segno.consts import ERROR_MAPPING

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)


@cli.command()
@click.option("--text", required=True)
@click.option("--background", type=ImageParamType(), required=True)
@click.option("--target", type=click.Path(dir_okay=False), required=True)
@click.option("--scale", type=click.IntRange(min=3, max=15), default=9)
@click.option("--error", type=click.Choice(list(ERROR_MAPPING)))
def generate_qr_code(text, background, target, scale, error):
    write_artistic(
        segno.make(text, error=error),
        background.filename,
        target,
        scale=scale,
    )