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

def load_amused_model():
    # pipeline = DiffusionPipeline.from_pretrained("Bakanayatsu/ponyDiffusion-V6-XL-Turbo-DPO") 
    # AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo"
    return AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")

# Generate image from prompt using AmusedPipeline
def generate_image(prompt):
    try:
        pipe = load_amused_model()
        generator = torch.Generator().manual_seed(8)  # Create a generator for reproducibility
        image = pipe(prompt, generator=generator).images[0]  # Generate image from prompt
        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()