app.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline, AutoPipelineForText2Image
import base64
from io import BytesIO



def text_to_image_model():
    # AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")
    # AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
     return AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")

# Generate image from prompt using AmusedPipeline
def generate_image(prompt):
    try:
        pipe = text_to_image_model()
        image = pipe(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()