Update app.py

app.py

@@ -8,23 +8,22 @@
 #     if text:
 #         out = pipe(text)
 #         st.write(out)
+
 import streamlit as st
-from transformers import AutoModelForImageGeneration, AutoTokenizer
-from PIL import Image
-import torch
+from diffusers import DiffusionPipeline
 
-model_name = "stabilityai/stable-diffusion-2-1" # loading the model
-model = AutoModelForImageGeneration.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load the stable-diffusion-2-1 model from diffusers
+pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
 
-text_input = st.text_area('Enter some text') #getting user text
+# Get input text from the user
+text_input = st.text_area('Enter some text')
 
+# Add a submit button
 if st.button('Generate Image'):
     if text_input:
-        input_ids = tokenizer(text_input, return_tensors="pt").input_ids #input token
-        with torch.no_grad(): #img generation
-            output = model.generate(input_ids)
-            
-        image = Image.fromarray(output[0].numpy().astype('uint8')) # certing tensor to imgonv
+        # Generate the image using the model
+        generated_image = pipeline(text_input)
+
+        # Display the generated image
+        st.image(generated_image, caption='Generated Image', use_column_width=True)
 
-        st.image(image, caption='Generated Image', use_column_width=True) #display output