|
|
import base64 |
|
|
from dataclasses import dataclass |
|
|
from io import BytesIO |
|
|
from pathlib import Path |
|
|
from typing import Literal, cast |
|
|
|
|
|
import gradio as gr |
|
|
import jinja2 |
|
|
from openai import OpenAI |
|
|
from PIL import Image |
|
|
from pydantic import BaseModel |
|
|
|
|
|
client = OpenAI() |
|
|
|
|
|
TEMPLATES_DIR = Path(__file__).resolve().parent / "templates" |
|
|
jinja_env = jinja2.Environment(loader=jinja2.FileSystemLoader(str(TEMPLATES_DIR))) |
|
|
|
|
|
SYSTEM_PROMPT = "You are expert prompt engineer" |
|
|
|
|
|
StyleName = Literal[ |
|
|
"General", |
|
|
"Fashion", |
|
|
"Emotional Lifestyle", |
|
|
"Extreme Sports", |
|
|
"Captivating", |
|
|
"Image Replication", |
|
|
"Red Bar Lighting", |
|
|
"Teal Noir", |
|
|
] |
|
|
|
|
|
|
|
|
@dataclass(frozen=True) |
|
|
class StyleDefinition: |
|
|
name: StyleName |
|
|
template_filename: str |
|
|
info: str |
|
|
|
|
|
|
|
|
STYLE_DEFINITIONS: dict[StyleName, StyleDefinition] = { |
|
|
"General": StyleDefinition( |
|
|
name="General", |
|
|
template_filename="general_prompt.jinja", |
|
|
info="Versatile, balanced storytelling with cinematic detail for most scenarios.", |
|
|
), |
|
|
"Fashion": StyleDefinition( |
|
|
name="Fashion", |
|
|
template_filename="fashion_prompt.jinja", |
|
|
info="Editorial fashion aesthetic highlighting garments, styling, and runway polish.", |
|
|
), |
|
|
"Emotional Lifestyle": StyleDefinition( |
|
|
name="Emotional Lifestyle", |
|
|
template_filename="emotional_lifestyle_prompt.jinja", |
|
|
info="Warm, candid lifestyle imagery that focuses on mood, relationships, and feelings.", |
|
|
), |
|
|
"Extreme Sports": StyleDefinition( |
|
|
name="Extreme Sports", |
|
|
template_filename="extreme_sports_prompt.jinja", |
|
|
info="High-adrenaline action shots that emphasize energy, motion, and athletic feats.", |
|
|
), |
|
|
"Captivating": StyleDefinition( |
|
|
name="Captivating", |
|
|
template_filename="captivating_prompt.jinja", |
|
|
info="Visually striking compositions with dramatic flair and memorable storytelling.", |
|
|
), |
|
|
"Image Replication": StyleDefinition( |
|
|
name="Image Replication", |
|
|
template_filename="image_replication_prompt.jinja", |
|
|
info=( |
|
|
"Mimic the reference image's composition, lighting, and styling exactly while" |
|
|
" inserting the user or their face in place of the original subject. Eg. If the reference image is a music album cover, the user's face will be embedded in the album cover." |
|
|
), |
|
|
), |
|
|
"Red Bar Lighting": StyleDefinition( |
|
|
name="Red Bar Lighting", |
|
|
template_filename="red_bar_lighting_prompt.jinja", |
|
|
info="Red bar lighting style for image generation.", |
|
|
), |
|
|
"Teal Noir": StyleDefinition( |
|
|
name="Teal Noir", |
|
|
template_filename="teal_noir_prompt.jinja", |
|
|
info="Teal noir style for image generation.", |
|
|
) |
|
|
} |
|
|
|
|
|
PROMPT_TEMPLATES = { |
|
|
style: jinja_env.get_template(config.template_filename) |
|
|
for style, config in STYLE_DEFINITIONS.items() |
|
|
} |
|
|
|
|
|
DEFAULT_STYLE: StyleName = "General" |
|
|
STYLE_CHOICES: tuple[StyleName, ...] = tuple(STYLE_DEFINITIONS.keys()) |
|
|
|
|
|
STYLE_INFORMATION_BLOCK = "\n".join( |
|
|
f"- {style}: {config.info}" for style, config in STYLE_DEFINITIONS.items() |
|
|
) |
|
|
|
|
|
|
|
|
class StyleSelectionResponse(BaseModel): |
|
|
style: StyleName |
|
|
|
|
|
|
|
|
def process_prompt(user_image, reference_image, target_label: str, user_prompt: str, style: StyleName) -> str: |
|
|
user_image_url = None |
|
|
reference_image_url = None |
|
|
|
|
|
if user_image is not None: |
|
|
buffer = BytesIO() |
|
|
user_image.convert("RGB").save(buffer, format="JPEG", quality=90) |
|
|
b64_image = base64.b64encode(buffer.getvalue()).decode("utf-8") |
|
|
user_image_url = f"data:image/jpeg;base64,{b64_image}" |
|
|
|
|
|
if reference_image is not None: |
|
|
buffer = BytesIO() |
|
|
reference_image.convert("RGB").save(buffer, format="JPEG", quality=90) |
|
|
b64_reference_image = base64.b64encode(buffer.getvalue()).decode("utf-8") |
|
|
reference_image_url = f"data:image/jpeg;base64,{b64_reference_image}" |
|
|
|
|
|
try: |
|
|
template = PROMPT_TEMPLATES[style] |
|
|
except KeyError as error: |
|
|
raise ValueError(f"Unsupported style: {style}") from error |
|
|
|
|
|
user_content = template.render(user_prompt=user_prompt) |
|
|
|
|
|
content = [{"type": "input_text", "text": user_content}] |
|
|
|
|
|
if user_image_url is not None: |
|
|
content.append({"type": "input_image", "image_url": user_image_url}) |
|
|
if reference_image_url is not None: |
|
|
content.append({"type": "input_image", "image_url": reference_image_url}) |
|
|
|
|
|
response = client.responses.create( |
|
|
model="gpt-5", |
|
|
reasoning={"effort": "minimal"}, |
|
|
input=[ |
|
|
{ |
|
|
"role": "system", |
|
|
"content": SYSTEM_PROMPT, |
|
|
}, |
|
|
{ |
|
|
"role": "user", |
|
|
"content": content, |
|
|
} |
|
|
], |
|
|
) |
|
|
return f"{response.output_text} {target_label.strip()}" |
|
|
|
|
|
|
|
|
def recommend_style(user_prompt: str, reference_image: Image.Image | None) -> StyleSelectionResponse: |
|
|
if reference_image is not None: |
|
|
buffer = BytesIO() |
|
|
reference_image.convert("RGB").save(buffer, format="JPEG", quality=90) |
|
|
b64_reference_image = base64.b64encode(buffer.getvalue()).decode("utf-8") |
|
|
reference_image_url = f"data:image/jpeg;base64,{b64_reference_image}" |
|
|
else: |
|
|
reference_image_url = None |
|
|
|
|
|
|
|
|
user_prompt = f"""You are an art director who must pick the most fitting style name for a user's prompt. |
|
|
Consider the available styles and choose the single best option. User has provided the reference image. |
|
|
|
|
|
Style Guide: |
|
|
{STYLE_INFORMATION_BLOCK} |
|
|
|
|
|
User Prompt: |
|
|
{user_prompt} |
|
|
""" |
|
|
content = [{"type": "input_text", "text": user_prompt}] |
|
|
if reference_image_url is not None: |
|
|
content.append({ |
|
|
"type": "input_image", "image_url": reference_image_url |
|
|
}) |
|
|
completion = client.responses.parse( |
|
|
model="gpt-5-mini", |
|
|
reasoning={"effort": "low"}, |
|
|
input=[{ |
|
|
"role": "user", |
|
|
"content": content, |
|
|
}], |
|
|
text_format=StyleSelectionResponse, |
|
|
) |
|
|
return completion.output_parsed.style |
|
|
|
|
|
|
|
|
def handle_auto_style_toggle(auto_enabled: bool) -> dict[str, object]: |
|
|
return gr.update(interactive=not auto_enabled) |
|
|
|
|
|
|
|
|
def generate_prompt_handler( |
|
|
user_image, |
|
|
reference_image, |
|
|
target_label: str, |
|
|
user_prompt: str, |
|
|
current_style: str | None, |
|
|
auto_style_enabled: bool, |
|
|
): |
|
|
|
|
|
if auto_style_enabled: |
|
|
current_style = recommend_style(user_prompt, reference_image) |
|
|
|
|
|
prompt_text = process_prompt( |
|
|
user_image=user_image, |
|
|
reference_image=reference_image, |
|
|
target_label=target_label, |
|
|
user_prompt=user_prompt, |
|
|
style=current_style, |
|
|
) |
|
|
display_text = f"Selected style: {current_style}\n\n{prompt_text}" |
|
|
return display_text, gr.update(value=current_style, interactive=False) |
|
|
|
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
user_image = gr.Image( |
|
|
label="Upload user photo", |
|
|
type="pil" |
|
|
) |
|
|
reference_image = gr.Image( |
|
|
label="Optional: Upload reference image (Eg. movie poster, music album cover, etc.)", |
|
|
type="pil", |
|
|
) |
|
|
target_label = gr.Textbox( |
|
|
label="Enter target label", |
|
|
placeholder="SMRA", |
|
|
) |
|
|
user_prompt = gr.Textbox( |
|
|
label="Enter your prompt", |
|
|
placeholder="picture of me while sitting in a chair in the ocean", |
|
|
lines=4, |
|
|
) |
|
|
style_dropdown = gr.Dropdown( |
|
|
choices=list(STYLE_CHOICES), |
|
|
value=DEFAULT_STYLE, |
|
|
label="Style Selection", |
|
|
info="Choose the visual style for your enhanced prompt", |
|
|
interactive=False, |
|
|
) |
|
|
auto_style_checkbox = gr.Checkbox( |
|
|
label="Auto-select best style", |
|
|
value=True, |
|
|
) |
|
|
generate_button = gr.Button("Generate Prompt") |
|
|
with gr.Column(): |
|
|
prompt_output = gr.Textbox( |
|
|
label="Style Prompt", |
|
|
lines=20, |
|
|
) |
|
|
|
|
|
generate_button.click( |
|
|
generate_prompt_handler, |
|
|
inputs=[ |
|
|
user_image, |
|
|
reference_image, |
|
|
target_label, |
|
|
user_prompt, |
|
|
style_dropdown, |
|
|
auto_style_checkbox, |
|
|
], |
|
|
outputs=[prompt_output, style_dropdown], |
|
|
) |
|
|
auto_style_checkbox.change( |
|
|
handle_auto_style_toggle, |
|
|
inputs=[auto_style_checkbox], |
|
|
outputs=[style_dropdown], |
|
|
) |
|
|
|
|
|
demo.launch() |
|
|
|
