app.py

import os, sys, json, re, time, base64, random, requests, shutil
import gradio as gr
from datetime import datetime
import urllib.parse
from groq import Groq
from exif import Image
from PIL import Image as PILImage, ExifTags as PILExifTags
import colorsys

MAX_SEED = 9999
MAX_IMAGE_SIZE = 2048

GROQ_APIKEY_PROMPTENHANCE = os.environ["GROQ_APIKEY_PROMPTENHANCE"]


# delete cache folder if exist and create new
CACHE_DIR = os.path.join(os.path.dirname(__file__), "cache")
IMAGE_DIR = os.path.join(CACHE_DIR, "images")
if os.path.exists(CACHE_DIR):
    shutil.rmtree(CACHE_DIR)
os.makedirs(CACHE_DIR)
os.makedirs(IMAGE_DIR)

RES = os.path.join(os.path.dirname(__file__), "_res")

gr.set_static_paths(paths=["_res/assets/"])

custom_css = RES + "/_custom.css"
custom_js = RES + "/_custom.js"

title = "Bilder Builder"

theme = gr.themes.Soft(
    primary_hue="cyan",
    radius_size="sm",
)


def read_exif(image_path):
    with open(image_path, "rb") as src:
        img = Image(src)
        img_comment = json.loads(img.user_comment)

        # checking if the key exists before removing
        if "concept" in img_comment:
            img_comment.pop("concept")

        return img_comment


def read_image_exfi_data(image_path):
    print("Imagepath:", image_path)
    img_exif_make, img_exif_comment = read_exif(image_path)
    return None, image_path, img_exif_comment


def groq_enhance_process(Prompt=""):
    client = Groq(api_key=GROQ_APIKEY_PROMPTENHANCE)
    Prompt = "random prompt" if Prompt == "" else Prompt
    SYSTEMPROMPT = os.path.join(RES, "groq_systemmessage_prompt_enhance.json")
    with open(SYSTEMPROMPT, "r") as f:
        SYSTEMPROMPT = json.load(f)

    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=[SYSTEMPROMPT, {"role": "user", "content": Prompt}],
        temperature=0.8,
        max_tokens=8000,
        top_p=0.9,
        stream=False,
        stop=None,
    )

    if completion.choices[0].message.content != "":
        enhanced_prompt = completion.choices[0].message.content
        enhanced_prompt = re.sub(r"[\.\"]+", "", enhanced_prompt)

    return enhanced_prompt


def image_get_size(image_path):
    img = PILImage.open(image_path)
    print("Image size:", img.size)
    width, height = img.size
    return width, height


def image_get_dominant_color(image_path):
    img = PILImage.open(image_path)
    img = img.convert("RGB")
    img = img.resize((100, 100), resample=0)
    pixels = list(img.getdata())

    # Erzeuge eine Liste mit den Häufigkeiten der Farben
    colors = []
    for pixel in pixels:
        r, g, b = pixel
        h, s, v = colorsys.rgb_to_hsv(r / 255, g / 255, b / 255)
        if v > 0.5:  # Filteriere hellere Farben aus
            continue
        if v > 0.99:  # Filteriere Weiß aus
            continue
        colors.append((h, s, v))

    # Ermittle die dominante Farbe
    dominant_color = max(colors, key=lambda x: x[2])
    dominant_color_rgb = colorsys.hsv_to_rgb(dominant_color[0], dominant_color[1], dominant_color[2])
    dominant_color_rgb = [int(c * 255) for c in dominant_color_rgb]
    dominant_color_rgb = f"rgb({dominant_color_rgb[0]}, {dominant_color_rgb[1]}, {dominant_color_rgb[2]})"
    print(dominant_color_rgb)

    return dominant_color_rgb


def process(Prompt, image_width, image_height, image_seed, randomize_seed):
    used_seed = random.randint(0, MAX_SEED) if image_seed == 0 or randomize_seed else image_seed
    if "test image" in Prompt.lower() or Prompt == "":
        print("Prompt:", Prompt)
        TESTIMAGES = os.path.join(RES, "testimages")
        random_num = random.randint(1, 10)
        match = re.match(r"test image (\d+)", Prompt)
        if match and 1 <= int(match.group(1)) <= 10:
            random_num = match.group(1)
            print("\n\nrandom: ", random_num)

        Prompt = "test image " + str(random_num) if Prompt == "" else Prompt
        filename = str(random_num) + ".png"
        print("Filename:", filename)
        file_path = os.path.join(TESTIMAGES, filename)

        img_exif_comment = read_exif(file_path)
        img_dominant_color = image_get_dominant_color(file_path)
        img_width, img_height = image_get_size(file_path)
        print("ImageSize: " + str(img_width) + "x" + str(img_height) + "px")

        time.sleep(1)

        return ({"value": file_path, "__type__": "update"}, {"value": None, "visible": False, "__type__": "update"}, img_exif_comment, {"visible": True, "__type__": "update"}, {"value": Prompt, "visible": True, "__type__": "update"}, img_width, img_height, used_seed, {"value": file_path, "visible": True, "__type__": "update"}, img_dominant_color)
    else:
        used_seed = random.randint(0, MAX_SEED) if image_seed == 0 or randomize_seed else image_seed
        used_model = "turbo" if int(image_width) > 1024 or int(image_height) > 1024 else "flux"  # turbo, flux

        timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        filename_prompt = re.sub(r"[^\w\s-]", "", Prompt).strip().replace(" ", "_")
        filename = timestamp + "_" + filename_prompt[:100] + ".png"
        # file_path = os.path.join(filename_dir, filename)
        file_path = os.path.join(IMAGE_DIR, filename)

        encode_prompt = urllib.parse.quote(Prompt)
        request_url = f"https://image.pollinations.ai/prompt/{encode_prompt}?model={used_model}&width={image_width}&height={image_height}&nologo=true&enhance=false&nofeed=true&seed={used_seed}"
        print(request_url)
        response = requests.get(request_url)
        if response.status_code == 200:
            with open(file_path, "wb") as f:
                f.write(response.content)

            img_exif_comment = read_exif(file_path)
            img_dominant_color = image_get_dominant_color(file_path)
            img_width, img_height = image_get_size(file_path)

            return ({"value": file_path, "__type__": "update"}, {"value": None, "visible": False, "__type__": "update"}, img_exif_comment, {"visible": True, "__type__": "update"}, {"value": Prompt, "visible": True, "__type__": "update"}, img_width, img_height, used_seed, {"value": file_path, "visible": True, "__type__": "update"}, img_dominant_color)


with gr.Blocks(theme=theme, css=custom_css, js=custom_js, title=title) as demo:
    with gr.Tab("Image Generator"):
        with gr.Row():
            with gr.Column(min_width=420, scale=1):
                with gr.Row():
                    placeholder_text = "[???] Generiert dir einen zufälligen Prompt.\n[STERN] optimiert deinen eignen Prompt.\n[RUN] generiert dein Bild."
                    text_prompt = gr.Textbox(label="Prompt", show_label=False, lines=12, max_lines=18, placeholder=placeholder_text, elem_classes="prompt-input", autofocus=True)
                with gr.Row():
                    random_prompt_button = gr.Button("", variant="secondary", elem_id="random_prompt_btn", elem_classes="random-prompt-btn", icon="_res/assets/star_light_48.png")
                    enhance_prompt_button = gr.Button("", variant="secondary", elem_id="enhance_prompt_btn", elem_classes="enhance-prompt-btn", icon="_res/assets/star_light_48.png")
                    run_button = gr.Button("Erstellen", variant="primary", elem_id="run_btn", elem_classes="run-btn")
                with gr.Row(elem_classes="image_size_selctor_wrapper"):
                    with gr.Column(scale=1):
                        with gr.Row():
                            # image_width = gr.Textbox(label="Breite", minimum=256, maximum=MAX_IMAGE_SIZE, value=1024, step=128, elem_id="image_width_selector", elem_classes="image-width-selector", scale=1)
                            image_width = gr.Textbox(label="Breite", value=1024, lines=1, max_lines=1, elem_id="image_width_selector", elem_classes="image-width-selector", scale=1)
                            # image_height = gr.Number(label="Höhe", minimum=256, maximum=MAX_IMAGE_SIZE, value=683, step=128, elem_id="image_height_selector", elem_classes="image-height-selector", scale=1)
                            image_height = gr.Textbox(label="Höhe", value=683, lines=1, max_lines=1, elem_id="image_height_selector", elem_classes="image-height-selector", scale=1)
                        with gr.Row():
                            image_ratio_buttons = gr.Radio(["16:9", "4:3", "3:2", "1:1"], value="3:2", label="Querformat", show_label=True, info="Die gängigsten Seitenverhältnissen.", interactive=True, elem_id="image_ratio_buttons", elem_classes="image-ratio-buttons", container=True, scale=2)
                            switch_width_height = gr.Button("", size="sm", elem_id="switch_width_height", elem_classes="switch-ratio-btn", variant="primary", scale=1)
                        with gr.Column():
                            image_seed = gr.Slider(label="Seed", info="Jeder Seed generiert ein anderes Bild mit dem selben Prompt", minimum=0, step=1, value=42, maximum=MAX_SEED)
                            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

            with gr.Column(min_width=600, scale=3):
                with gr.Row():
                    with gr.Column(scale=1):
                        with gr.Row():
                            output_image = gr.Image(show_label=False, height=720, min_width=320, scale=3, elem_id="output_image", type="filepath")
                            with gr.Column(scale=1, visible=False, elem_classes="image-info-wrapper") as image_info_wrapper:
                                # with gr.Group(visible=False) as image_info_wrapper:
                                with gr.Group():
                                    image_informations = gr.Markdown("""## Bildinformationen""", visible=True)
                                # textbox_your_prompt = gr.Textbox(label="Dein Prmopt", lines=2, max_lines=4, interactive=False, show_copy_button=True)
                                with gr.Row(elem_classes="img-size-wrapper"):
                                    image_info_tb_width = gr.Textbox(label="Breite", lines=1, max_lines=1, interactive=False, show_copy_button=True, elem_classes="image-info-tb-width")
                                    image_info_tb_height = gr.Textbox(label="Höhe", lines=1, max_lines=1, interactive=False, show_copy_button=True, elem_classes="image-info-tb-height")
                                with gr.Row(elem_classes="img-seed-wrapper"):
                                    image_info_tb_seed = gr.Textbox(label="Seed", lines=1, max_lines=1, interactive=False, show_copy_button=True, elem_classes="image-info-tb-seed")
                                image_info_tb_prompt = gr.Textbox("Bild Prompt", lines=4, max_lines=8, interactive=False, show_copy_button=True, visible=True)
                                image_download_button = gr.DownloadButton("Bild herunterladen", value=None, elem_classes="download-button", variant="primary", visible=False)

                output_url = gr.Textbox(label="Output URL", show_label=True, interactive=False, visible=False)
                outpu_image_comment = gr.Json(visible=False)
                output_dominant_image_color = gr.Textbox(show_label=False, elem_id="dominant_image_color", visible=True, elem_classes="output-dominant-image-color")
   

        def switch_image_size_values(image_width, image_height):
            return image_height, image_width

        # switch_width_height.click(fn=switch_image_size_values, inputs=[image_width, image_height], outputs=[image_width, image_height])

        run_button.click(fn=process, inputs=[text_prompt, image_width, image_height, image_seed, randomize_seed], outputs=[output_image, output_url, outpu_image_comment, image_informations, image_info_tb_prompt, image_info_tb_width, image_info_tb_height, image_info_tb_seed, image_download_button, output_dominant_image_color])

        enhance_prompt_button.click(fn=groq_enhance_process, inputs=[text_prompt], outputs=[text_prompt])
        random_prompt_button.click(fn=groq_enhance_process, inputs=None, outputs=[text_prompt])
        

demo.launch()