import gradio as gr
import spaces

#gr.load("models/sergon19/green_bg_LoRa10").launch()

import torch
from diffusers import DiffusionPipeline, AutoencoderKL

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)
pipe.load_lora_weights("sergon19/green_bg_LoRa10")
_ = pipe.to("cuda")

@spaces.GPU
def generate_image(prompt):
    image = pipe(prompt=prompt, num_inference_steps=25).images[0]
    return image

# Adding examples to the Gradio interface
examples = [
    ["A serene landscape with mountains and a river in sgc style"],
    ["A futuristic city skyline at night in sgc style"],
    ["a busy wallstreet office in sgc style"],
    ["A beach scene with a pirate ship visible in the ocean in sgc style"],
    ["A vibrant underwater scene with coral reefs and fish in sgc style"],
    ["A magical forest with glowing plants and creatures in sgc style"],
]

iface = gr.Interface(fn=generate_image, inputs="text", outputs="image", examples=examples)
iface.launch()