import gradio as gr
from diffusers import DiffusionPipeline
import torch
import base64
import io

pipe = DiffusionPipeline.from_pretrained("Nihirc/Prompt2MedImage")
if torch.cuda.is_available():
    pipe = pipe.to("cuda")

def generate_image(prompt):
    image = pipe(prompt).images[0]
    
    # Convert to base64 for API compatibility
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    
    return img_str  # Return base64 string directly

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Textbox(label="Base64 Image"),  # Changed from Image to Textbox
    title="Medical Image Generator"
)

iface.launch()