File size: 1,065 Bytes
f466dd9
 
a9b8939
9e09422
 
f466dd9
a9b8939
 
6d1d03a
f466dd9
 
a9b8939
f466dd9
 
 
 
 
a9b8939
 
f466dd9
 
a9b8939
 
f466dd9
bf161ed
9e09422
 
 
 
f466dd9
 
 
 
9e09422
f466dd9
 
 
a9b8939
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
import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image
import base64
from io import BytesIO

# Load the model once outside of the function
model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")

def generate_image(prompt):
    try:
        image = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
        return image, None
    except Exception as e:
        return None, str(e)

def inference(prompt):
    print(f"Received prompt: {prompt}")
    # Debugging statement
    image, error = generate_image(prompt)
    if error:
        print(f"Error generating image: {error}")
        # Debugging statement
        return "Error: " + error
    
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return img_str

gradio_interface = gr.Interface(
    fn=inference,
    inputs="text",
    outputs="text"  # Change output to text to return base64 string
)

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