gen_tab.py

import gradio as gr
from src.utils import add_to_prompt
from src.rep_api import generate_image_replicate, model_dict
import os

# The dictionary data
prompt_dict = {
    "Character": ["Asian girl with black hair", "A man with blond hair", "A Cat girl anime character with purple hair", "A Green Alien with big black eyes"],
    "Clothes": ["Wearing a blue jacket", "Wearing a black business suit", "Wearing a purple jumpsuit", "Wearing shorts and a white T-shirt"],
    "Pose": ["Close up portrait", "Standing doing a peace sign", "Folding arms", "holding a phone"],
    "Style": ["Simple white background", "Fashion runway", "Inside a business conference", "Inside a spaceship"],
}


def create_gen_tab():
    with gr.TabItem("Image Generator"):
        with gr.Row():
            with gr.Column():
                inp = gr.Textbox(label="Prompt")
                with gr.Accordion("extra", open=False):
                    replicate_api_key = gr.Textbox(
                        label="Replicate API Key",
                        info="API key for Replicate",
                        value=os.environ.get("REPLICATE_API_TOKEN", ""),
                        type="password",
                    )
                    model_type= gr.Dropdown(["dev","schnell"],value="dev",label="flux model base", info="dev - higher quailty but longer \n schnell lower quaility but faster")


                btn = gr.Button("Generate")
            with gr.Column():
                aspect_ratio = gr.Dropdown(list(["1:1","16:9","9:16","5:3"]),value="1:1", label="Aspect Ratio", info="Aspect Ratio")
                style_mode = gr.Dropdown(list(model_dict.keys()),label="Style lore",value="Base",info="style model")
                style_strength = gr.Slider(0,2,value=1,label="Style Strength")
                api_path = gr.Textbox(label="Other_style",info="Style lora safytensor route goes here or replicate link goes here",value=None)
                
                
        with gr.Accordion("Prompt Support", open=False):
            for key, values in prompt_dict.items():
                with gr.Row():
                    gr.Button(key,interactive=False)
                    for value in values:
                        gr.Button(value).click(add_to_prompt, inputs=[inp, gr.Textbox(value,visible=False)], outputs=inp)

        with gr.Row():
            gen_out = gr.Image(label="Generated Image",type="filepath")
        with gr.Accordion("Gallery", open=False):
            gen_gallery = gr.Gallery(label="Generated Images Gallery", type="filepath")

        btn.click(generate_image_replicate, inputs=[inp,style_mode,api_path,aspect_ratio,gen_gallery,model_type,style_strength], outputs=[gen_out,gen_gallery])