app.py

import gradio as gr
import requests
import base64
from PIL import Image
from io import BytesIO


def decode_base64_image(image_string):
    base64_image = base64.b64decode(image_string)
    buffer = BytesIO(base64_image)
    return Image.open(buffer)


def inference(prompt, guidance_scale, num_inference_steps):
    api_url = 'https://a02q342s5b.execute-api.us-east-2.amazonaws.com/reinvent-demo-inf2-sm-20231114'

    prompt_input_one = {
        "prompt": prompt,
        "parameters": {
            "num_inference_steps": num_inference_steps,
            "guidance_scale": guidance_scale,
            "seed": -1
        },
        "endpoint": "huggingface-pytorch-inference-neuronx-2023-11-14-21-22-10-388"
    }

    response_one = requests.post(api_url, json=prompt_input_one)

    if response_one.status_code == 200:
        result_one = response_one.json()
        return decode_base64_image(result_one["generated_images"][0])
    else:
        return None


# example_prompts = [
#     "Illustrate the concept of 'serendipity' in a single image.",
#     "Create a visual representation of the phrase 'whispers of the autumn wind.'",
#     "Design an image capturing the essence of 'timeless wonder' in a mystical forest setting.",
#     "Visualize the emotions evoked by the words 'bittersweet symphony' in a unique artwork.",
#     "Craft an image that conveys the juxtaposition of 'urban chaos and tranquil solitude' in a futuristic cityscape.",
#     "Albert Einstein smoking, half body wide angle professional portrait photo, clear and detailed eyes, "
#     "nature background at sunset, backlighting, fashion photography, centered, symmetrical, hasselblad helios 44-2 "
#     "58mm F2, by Annie Leibovitz and Ellen von Unwerth"
# ]


def app():
    return gr.Interface(inference,
                        [gr.Textbox(
                            label="Prompt",
                            info="Enter your prompt",
                            lines=3,
                            value="Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
                        ),
                            gr.Slider(2, 20, value=7.5, label="Guidance Scale"),
                            gr.Slider(1, 50, value=4, label="Inference steps")
                        ],

                        gr.Image(type="pil",
                                 height=512,
                                 width=512
                                 )
                        , allow_flagging='never', title='Gen Image',
                        examples=[
                            ["Self-portrait oil painting, a beautiful cyborg with golden hair, 8k", 7, 10],
                            ["Self-portrait oil painting, a beautiful cyborg with golden hair, 8k", 15, 20],
                            ["Design an image capturing the essence of 'timeless wonder' in a mystical forest setting.",
                             7, 20],
                            ["Visualize the emotions evoked by the words 'bittersweet symphony' in a unique artwork.",
                             15, 20],
                        ]
                        )


if __name__ == "__main__":
    app().launch()