Update app.py

app.py

@@ -1,43 +1,28 @@
 import streamlit as st
-from diffusers import StableDiffusionPipeline
-import torch
-from PIL import Image
-import io
-import base64
+from transformers import DiffusionPipeline
 
-# Model ID
-model_id = "stabilityai/stable-diffusion-2-1"
+# Load the Diffusion pipeline
+pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium")
 
-# Load the model from Hugging Face Hub
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+def generate_prompt(prompt_text):
+    # Generate response using the Diffusion model
+    response = pipeline(prompt=prompt_text, top_p=0.9, num_return_sequences=1)[0]['generated_text']
+    return response
 
-# Streamlit app title and description
-st.title("Stable Diffusion Image Generation")
-st.write("Generate images using the Stability AI's Stable Diffusion model.")
+# Streamlit app UI
+st.title('Diffusion Model Prompt Generator')
 
-# Input for the prompt
-prompt = st.text_input("Enter a prompt for the image:")
+# Input prompt from user
+prompt_input = st.text_area('Enter your prompt here:', height=100)
 
 # Generate button
-if st.button("Generate Image"):
-    if prompt:
-        with st.spinner('Generating image...'):
-            # Generate image
-            image = pipe(prompt).images[0]
-
-            # Convert image to displayable format
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
+if st.button('Generate'):
+    if prompt_input:
+        with st.spinner('Generating...'):
+            generated_text = generate_prompt(prompt_input)
+        st.success('Generation complete!')
+        st.text_area('Generated Text:', value=generated_text, height=200)
+    else:
+        st.warning('Please enter a prompt to generate.')
 
-            # Display the image
-            st.image(image, caption="Generated Image", use_column_width=True)
 
-            # Option to download the image
-            st.markdown(
-                f'<a href="data:image/png;base64,{img_str}" download="generated_image.png">Download Image</a>',
-                unsafe_allow_html=True,
-            )
-    else:
-        st.error("Please enter a prompt to generate an image.")