import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

def generate(prompt):
    base_model = "runwayml/stable-diffusion-v1-5"
    lora_model = "glif-loradex-trainer/djmenorobl_PS2_STYLE_IMAGE_GENERATOR"

    pipe = StableDiffusionPipeline.from_pretrained(
        base_model,
        torch_dtype=torch.float16
    ).to("cuda")

    pipe.load_lora_weights(lora_model)
    image = pipe(prompt).images[0]
    return image

gr.Interface(fn=generate, inputs="text", outputs="image", title="PS2 Style Image Generator").launch()