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

pipe = DiffusionPipeline.from_pretrained("segmind/tiny-sd")
pipe.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, num_samples):
    return pipe(
        prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        num_images_per_prompt=num_samples
    ).images


app = gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Text("", label="Negative Prompt"),
        gr.Slider(minimum=1, maximum=100, value=45, label="Number inference steps"),
        gr.Number(512, label='Image width'),
        gr.Number(512, label='Image height'),
        gr.Slider(minimum=1, maximum=10, value=1, label='Number samples'),
    ],
    outputs=gr.Gallery(),
    examples=[
        [
            "Colorful line shading by JON_JUAREZ a dark cave with toxic mushrooms",
            "",
            45,
            512,
            512,
            1
        ]
        
        
    ]
)
    
app.launch()