import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the model from your repository
model_id = "linnthit111/imagemm"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)


def generate_image(prompt):
    # ပုံဖန်တီးတဲ့အခါ safety_checker ကို ဖြုတ်ပါ
image = pipe(prompt).images[0]

    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Myanmar Text-to-Image",
    description="Generate images from text prompts using a custom Stable Diffusion model."
)

iface.launch()