Update app.py

app.py

@@ -1,32 +1,36 @@
 import streamlit as st
 import torch
 from diffusers import StableDiffusionPipeline
-import matplotlib.pyplot as plt
 
-st.title('Text-to-Image Generation using Stable Diffusion')
-
-# Load the diffusion model
 model_id1 = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
 pipe = pipe.to("cuda")
 
-# Define function to generate images
-def generate_image(prompt, negative_prompt, num_inference_steps=50, weight=640):
-    params = {'prompt': prompt, 'num_inference_steps': num_inference_steps, 'num_images_per_prompt': 2, 'height': int(1.2 * weight),
-            'weight': weight, 'negative_prompt': negative_prompt}
-
+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]
 
-# Create Streamlit interface
-st.sidebar.subheader("Enter Prompts")
-prompt = st.sidebar.text_area("Enter Prompt")
-negative_prompt = st.sidebar.text_area("Enter Negative Prompt")
-num_inference_steps = st.sidebar.slider("Number of Inference Steps", 1, 100, 50)
-weight = st.sidebar.slider("Image Width", 512, 640, 640)
 
-if st.sidebar.button("Generate Images"):
-    image1, image2 = generate_image(prompt, negative_prompt, num_inference_steps, weight)
+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)
+
 
-    st.image(image1, caption='Image 1', use_column_width=True)
-    st.image(image2, caption='Image 2', use_column_width=True)
+if __name__ == "__main__":
+    main()