import gradio as gr
import os
import spaces
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained("segmind/tiny-sd")
pipeline.load_lora_weights(
    "philipp-zettl/jon_juarez-lora",
    hf_token=os.environ.get('HF_TOKEN')
)
pipe.to('cuda')
@spaces.GPU
def generate(prompt, negative_prompt, num_inference_steps, width, height):
    return pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, width=width, height=height).images


app = gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Text("", label="Negative Prompt"),
        gr.Number(45, label="Number inference steps"),
        gr.Number(1024, label='image width'),
        gr.Number(1024, label='image height'),
    ],
    outputs=gr.Gallery(),
)
with app as demo:

    demo.examples = [
        "Colorful line shading by JON_JUAREZ a dark cave with toxic mushrooms",
    
    ]
    gr.Slider(label="seed")

    
demo.launch()