app.py

import gradio as gr
from openai import OpenAI
from os import environ

# Set the API key
Mykey = environ.get('OpenAIKEY')
client = OpenAI(api_key=Mykey)

MODEL = "gpt-4o-2024-05-13"

def translate_text(text, target_language):
    response = client.chat.completions.create(
        model=MODEL,  # Specify the GPT-4 model
        messages=[
            {"role": "system", "content": f"You are a translator to {target_language}."},
            {"role": "user", "content": text}
        ],
        max_tokens=100,
        temperature=0.7,
    )
    print(response)  # Debugging: print the response object
    translation = response.choices[0].message.content.strip()
    return translation

def refine_prompt(prompt):
    response = client.chat.completions.create(
        model=MODEL,  # Specify the GPT-4 model
        messages=[
            {
                "role": "system", 
                "content": (
                    "You are a prompt refiner for Stable Diffusion image generation. "
                    "Your goal is to enhance user prompts to make them more vivid, detailed, and suitable for generating high-quality images."
                    "Ensure the refined prompts include specific details about the scene, characters, environment, lighting, and artistic style."
                    "Avoid ambiguous or generic descriptions."
                    "All descriptions should be descriptive of one particular scene and should not have sequential steps. If there are more than one main object in the scene, you should display all the main objects in the same image"
                    "Here are some examples of good prompts:\n"
                    "1. 'A bowl filled with soaking wood ear mushrooms expanding in water. Next to it, a clean cutting board with a knife and freshly washed wood ear mushrooms ready to be sliced. A pot of boiling water on a stove, with steam rising.'\n"
                    "2. 'A bunch of fresh garlic sprouts being washed under running water in a sink. On the counter, some garlic sprouts are cut into neat sections on a cutting board with a knife.'\n"
                    "3. 'A mixing bowl with three cracked eggs, beaten and mixed with strips of wood ear mushrooms. A hand holding a spoon, mixing the ingredients thoroughly. A salt shaker nearby.'\n"
                    "Use these examples as a guide to refine the user's prompt."
                )
            },
            {"role": "user", "content": prompt}
        ],
        max_tokens=100,
        temperature=0.7,
    )
    print(response)  # Debugging: print the response object
    refined_prompt = response.choices[0].message.content.strip()
    return refined_prompt

def process(text, target_language):
    translated_text = translate_text(text, target_language)
    refined_prompt = refine_prompt(translated_text)
    return translated_text, refined_prompt

# Define the Gradio interface
iface = gr.Interface(
    fn=process,
    inputs=[
        gr.Textbox(lines=2, placeholder="待翻译文本", label="待翻译文本"),
        gr.Textbox(lines=1, placeholder="这里输出翻译目标语言", label="语言")
    ],
    outputs=[
        gr.Textbox(label="翻译结果"),
        gr.Textbox(label="精炼提示")
    ],
    title="提示翻译与润色工具",
    description="将文本翻译为指定语言，并优化提示语以便更好地使用Stable Diffusion生成图像。"
)

iface.launch(share=True)