wan2.1-Diffuser / app.py
rahul7star's picture
Update app.py
8ab0463 verified
raw
history blame
827 Bytes
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline
# Model path from Hugging Face
MODEL_NAME = "Evados/DiffSynth-Studio-Lora-Wan2.1-ComfyUI"
LORA_FILE = "Wan2.1-1.3b-lora-aesthetics-v1_new.safetensors"
def load_model():
pipe = StableDiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights(f"{MODEL_NAME}/{LORA_FILE}")
return pipe
pipe = load_model()
def generate_image(prompt):
image = pipe(prompt).images[0]
return image
# Gradio UI
iface = gr.Interface(
fn=generate_image,
inputs=gr.Textbox(label="Enter your prompt"),
outputs=gr.Image(label="Generated Image"),
title="Wan2.1 LoRA Image Generator",
description="Generate images using the Wan2.1 LoRA model. Enter a prompt to begin."
)
iface.launch()