import gradio as gr
from gradio_client import Client, handle_file

PROMT_GET_STYLE = 'Describe the artistic style and technique of this image in one sentence, DO NOT DESCRIBE THE PLOT. Start with: "Change the style of this picture to"'

def get_joy_caption(
    image_path: str,
    prompt: str = PROMT_GET_STYLE,
    temperature: float = 0.6,
    top_p: float = 0.9,
    max_new_tokens: int = 512,
    log_prompt: bool = True
) -> str:
    client = Client("fancyfeast/joy-caption-beta-one")
    result = client.predict(
        input_image=handle_file(image_path),
        prompt=prompt,
        temperature=temperature,
        top_p=top_p,
        max_new_tokens=max_new_tokens,
        log_prompt=log_prompt,
        api_name="/chat_joycaption"
    )
    return result

with gr.Blocks() as demo:
    gr.Markdown("## Flux Kontext Prompt Generator — Style Transfer")

    with gr.Row():
        image = gr.Image(type="filepath", label="Upload or Paste Image")
        prompt_out = gr.Textbox(label="Generated Prompt", lines=2, show_copy_button=True)
    
    with gr.Row():
        btn = gr.Button("Get Style")

    btn.click(fn=get_joy_caption, inputs=image, outputs=prompt_out)
    
if __name__ == "__main__":
    demo.launch()