joy-caption-cli/cli.py at master · inflamously/joy-caption-cli · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import click
from PIL import ImageFile

from genai.commands import genai
from quality.commands import quality_check
from segmentation.segmentation import segmentation, has_gen2seg_installed
from captions.clip.folder import folder as clip_caption_folder
from captions.joy.file import file as joy_file_command
from captions.joy.folder import folder as joy_folder_command
from captions.joy.profiling import caption_profile_image as joy_caption_profile_image
from initialization import setup_config


# File: cli.py
# Author: nflamously
# Original License: Apache License 2.0


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


@click.command()
def version():
    return "1.0"


@click.group("joycaption")
def caption():
    setup_config()


@click.group("clip")
def clip_caption():
    pass


if __name__ == "__main__":
    # Image edge case to handle
    ImageFile.LOAD_TRUNCATED_IMAGES = True

    # Joy Caption
    caption.add_command(joy_file_command)
    caption.add_command(joy_folder_command)
    caption.add_command(joy_caption_profile_image)

    # Stable Diffusion 1.5 -> CLIP
    clip_caption.add_command(clip_caption_folder)

    cli.add_command(quality_check)
    cli.add_command(clip_caption)
    cli.add_command(version)
    cli.add_command(caption)
    cli.add_command(genai)

    if has_gen2seg_installed():
        cli.add_command(segmentation)

    cli()