Update app.py

app.py

@@ -1,31 +1,26 @@
+import torch
 import streamlit as st
-from transformers import pipeline
+from diffusers import StableDiffusion3Pipeline
 
-# 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 = generator(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
-    return response
+# Load the model
+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 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__':
+    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()
 
 
+