import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

# Load the base model
base_model = "stabilityai/stable-diffusion-xl-base-1.0"
lora_model = "itsVilen/Mspaint"

pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

# Load the LoRA weights
pipe.load_lora_weights(lora_model)
pipe.to("cuda")

def generate_image(prompt):
    try:
        # Generate an image from the prompt
        image = pipe(prompt).images[0]
        return image
    except Exception as e:
        return str(e)

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
    outputs=gr.Image(type="pil"),
    title="Text to Image Generation",
    description="Enter a text prompt and generate an image using the itsVilen/Mspaint model.",
)

iface.launch()