import matplotlib.pyplot as plt
import gradio as gr

from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
import torch

model_id1 = "dreamlike-art/dreamlike-diffusion-1.0"
model_id2 = "stabilityai/stable-diffusion-xl-base-1.0"
model_id3 = "stabilityai/stable-diffusion-2"

pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
pipe = pipe.to("cuda")

def generate_image_interface(prompt, num_inference_steps, height, width):
    params = {
        'prompt': prompt,
        'num_inference_steps': num_inference_steps,
        'num_images_per_prompt': 2,
        'height': height,
        'width': width
    }
    
    img = pipe(**params).images  # Ensure the `pipe` call correctly matches the expected API
    
    num_images = len(img)
    if num_images > 1:
        fig, ax = plt.subplots(nrows=1, ncols=num_images, figsize=(15, 5))
        for i in range(num_images):
            ax[i].imshow(img[i])
            ax[i].axis('off')
    else:
        fig = plt.figure()
        plt.imshow(img[0])
        plt.axis('off')
        
    plt.tight_layout()
    plt.show()
    return fig

# Define the Gradio interface
inputs = [
    gr.Textbox(label="Enter your prompt"),
    gr.Slider(minimum=1, maximum=100, value=50, label="Number of Inference Steps"),
    gr.Slider(minimum=512, maximum=1024, value=768, label="Height"),
    gr.Slider(minimum=512, maximum=1024, value=768, label="Width")
]
outputs = gr.Plot()

demo = gr.Interface(fn=generate_image_interface, inputs=inputs, outputs=outputs)
demo.launch(share=True)