import torch
from diffusers import DiffusionPipeline
import gradio as gr

# Detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Load pipeline
pipe = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=dtype
)
pipe.to(device)

# Load LoRA weights (requires `peft` installed)
pipe.load_lora_weights("EliKet/train_text_to_img")

# Inference function
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio Interface
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Describe the image you want..."),
    outputs="image",
    title="🖼️ LoRA Text-to-Image Generator",
    description="Enter a prompt to generate an image using Stable Diffusion with LoRA (EliKet/train_text_to_img)."
)

# Launch app
demo.launch()