import gradio as gr
import requests
import json

# Define the function to process the input
def process_input(user_input):
    # Append "<mask>" to the user input
    input_with_mask = user_input + " <mask>"
    
    # Define the API endpoint and headers
    url = "https://44ts8edkp0.execute-api.us-east-1.amazonaws.com/PROD/prompt-generation"
    headers = {"Content-Type": "application/json"}
    
    # Create the data payload
    data = {"inputs": input_with_mask}
    
    # Make the POST request
    response = requests.post(url, headers=headers, data=json.dumps(data))
    
    # Parse the JSON response
    response_json = response.json()
    
    # Extract the "sequence" values from the response
    sequences = [item["sequence"] for item in response_json]
    
    # Join the sequences into a single string for display
    result = "\n".join(sequences)
    
    return result

# Create the Gradio interface
iface = gr.Interface(
    fn=process_input,
    inputs="text",
    outputs="text",
    title="Text-to-image prompt generator",
    description="Start writing your prompt and the the app will give you useful suggestions to fill and suggest different ways to write your prompt"
)

# Launch the app
if __name__ == "__main__":
    iface.launch()