import gradio as gr
from transformers import pipeline
import os
import io
import IPython.display
from PIL import Image
import base64

hf_api_key = "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"
import requests, json

def get_completion(inputs, parameters=None, ENDPOINT_URL="https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"):
    headers = {
      "Authorization": f"Bearer {hf_api_key}",
      "Content-Type": "application/json"
    }
    data = { "inputs": inputs }
    if parameters is not None:
        data.update({"parameters": parameters})
    response = requests.request("POST",
                                ENDPOINT_URL,
                                headers=headers,
                                data=json.dumps(data))
    return response.content

import gradio as gr

# A helper function to convert the PIL image to base64
# so you can send it to the API
def base64_to_pil(img_base64):
    base64_decoded = base64.b64decode(img_base64)
    byte_stream = io.BytesIO(base64_decoded)
    pil_image = Image.open(byte_stream)
    return pil_image

def generate(prompt):
    output = get_completion(prompt)
    result_image = Image.open(io.BytesIO(output))
    return result_image

gr.close_all()
demo = gr.Interface(
    fn=generate,
    inputs=[gr.Textbox(label="Your prompt")],
    outputs=[gr.Image(label="Result")],
    title="Image Generation with Stable Diffusion",
    description="Generate any image with Stable Diffusion",
    allow_flagging="never",
)

demo.launch(inline = False)