import os
import io
from PIL import Image
import base64
from dotenv import load_dotenv, find_dotenv
import gradio as gr
import requests,json

_ = load_dotenv(find_dotenv())  # read local .env file

hf_api_key = os.environ['HF_API_KEY']

# Adjusted Helper function
def get_completion(inputs, parameters=None, ENDPOINT_URL=os.environ['HF_API_TTI_BASE']):
    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.post(ENDPOINT_URL, headers=headers, json=data)

    if 'application/json' in response.headers.get('Content-Type'):
        return response.json()  # If response is JSON
    else:
        # If response is not JSON, handle as binary (image data)
        return base64.b64encode(response.content).decode('utf-8')  # Convert binary image to base64

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

#Updated generate function to handle base64 image string
def generate(prompt):
    output = get_completion(prompt)
    # Assuming output is now a base64 encoded string of the image
    pil_image = base64_to_pil(output)  # Convert base64 string to PIL Image
    return pil_image

# Rest of your Gradio setup remains the same
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(share=True)