Update app.py

app.py

@@ -1,26 +1,27 @@
-import torch
 import streamlit as st
-from diffusers import StableDiffusion3Pipeline
+from diffusers import DiffusionPipeline
 
-# Retry mechanism for loading the model
-def retry_load_model():
+# Load the Diffusion pipeline
+@st.cache(allow_output_mutation=True)
+def load_diffusion_pipeline():
     try:
-        return StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
-    except EnvironmentError as e:
+        pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium")
+        return pipeline
+    except Exception as e:
         st.error(f"Error loading model: {e}")
-        st.warning("Retrying to load the model...")
-        retry_load_model()
-
-# Load the Diffusion pipeline
-pipeline = retry_load_model()
 
-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 generate_response(prompt_text, pipeline):
+    try:
+        response = pipeline(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
+        return response
+    except Exception as e:
+        st.error(f"Error generating response: {e}")
 
 def main():
-    st.title('Diffusion Model Prompt Generator')
+    st.title('Hugging Face Diffusion Model')
+
+    # Load the model
+    pipeline = load_diffusion_pipeline()
 
     # Text input for the prompt
     prompt_text = st.text_area("Enter your prompt here:", height=200)
@@ -29,7 +30,7 @@ def main():
     if st.button("Generate"):
         if prompt_text:
             with st.spinner('Generating...'):
-                generated_text = generate_prompt(prompt_text)
+                generated_text = generate_response(prompt_text, pipeline)
             st.success('Generation complete!')
             st.text_area('Generated Text:', value=generated_text, height=400)
         else: