File size: 1,835 Bytes
532c92e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
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)
|