import gradio as gr
from diffusers import DiffusionPipeline

# Load the DiffusionPipeline and LORA weights
pipeline = DiffusionPipeline.from_pretrained("stablediffusionapi/juggernaut-xl-v5")
pipeline.load_lora_weights("Kvikontent/kviimager2.0")

# Define the function to handle user input and generate the image
def generate_image(prompt):
    image = pipeline(prompt)
    return image

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="KVIImager 2.0 Demo",
    description="Enter a prompt to generate an image"
)

# Add examples for the user input
examples = [
    "Sunset over the ocean",
    "Starry night sky"
]

iface.set_config(
    examples=examples
)

# Launch the Gradio interface
iface.launch()