import streamlit as st
import torch
from diffusers import StableDiffusionPipeline

model_id1 = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
pipe = pipe.to("cuda")

def generate_image(prompt, negative_prompt, num_inference_steps=50, width=640):
    params = {
        'prompt': prompt,
        'num_inference_steps': num_inference_steps,
        'num_images_per_prompt': 2,
        'height': int(1.2 * width),
        'width': width,
        'negative_prompt': negative_prompt
    }
    img = pipe(**params).images
    return img[0], img[1]


def main():
    st.title("Diffuser Image Generator")
    prompt = st.text_input("Enter the prompt:")
    negative_prompt = st.text_input("Enter the negative prompt:")
    num_inference_steps = st.slider("Number of inference steps", 1, 100, 50)
    width = st.slider("Width", 512, 640, 640)

    if st.button("Generate Image"):
        image1, image2 = generate_image(prompt, negative_prompt, num_inference_steps, width)
        st.image(image1, caption="Generated Image 1", use_column_width=True)
        st.image(image2, caption="Generated Image 2", use_column_width=True)


if __name__ == "__main__":
    main()