ma4389's picture
Update app.py
f97ce46 verified
raw
history blame contribute delete
903 Bytes
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()