Update app.py

app.py

@@ -1,28 +1,31 @@
 import streamlit as st
-from transformers import DiffusionPipeline
+from transformers import pipeline
 
-# Load the Diffusion pipeline
-pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium")
+# Load the Diffusion pipeline for text generation
+generator = pipeline("text-generation", model="stabilityai/stable-diffusion-3-medium")
 
 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']
+    response = generator(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
     return response
 
-# Streamlit app UI
-st.title('Diffusion Model Prompt Generator')
-
-# Input prompt from user
-prompt_input = st.text_area('Enter your prompt here:', height=100)
-
-# Generate button
-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.')
+def main():
+    st.title('Diffusion Model Prompt Generator')
+
+    # Text input for the prompt
+    prompt_text = st.text_area("Enter your prompt here:", height=200)
+
+    # Button to generate prompt
+    if st.button("Generate"):
+        if prompt_text:
+            with st.spinner('Generating...'):
+                generated_text = generate_prompt(prompt_text)
+            st.success('Generation complete!')
+            st.text_area('Generated Text:', value=generated_text, height=400)
+        else:
+            st.warning('Please enter a prompt.')
+
+if __name__ == '__main__':
+    main()