Update app.py

app.py

@@ -1,26 +1,31 @@
-import torch
 import streamlit as st
 from diffusers import StableDiffusion3Pipeline
 
-# Load the model
+# Load the Diffusion pipeline
 pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
-pipeline = pipeline.to("cuda")  # Move model to GPU if available
 
-# Streamlit UI
+def generate_prompt(prompt_text):
+    # Generate response using the Diffusion model
+    response = pipeline(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
+    return response
+
 def main():
-    st.title("Stable Diffusion 3 Medium Demo")
-    prompt = st.text_input("Enter your prompt:", "A cat holding a sign that says hello world")
-
-    if st.button("Generate Image"):
-        with st.spinner("Generating..."):
-            try:
-                image = pipeline(prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0).images[0]
-                st.image(image, caption="Generated Image", use_column_width=True)
-            except Exception as e:
-                st.error(f"Error: {e}")
-
-if __name__ == "__main__":
-    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()