Update app.py

app.py

@@ -2,11 +2,17 @@ import torch
 import streamlit as st
 from diffusers import StableDiffusion3Pipeline
 
-# Load the Diffusion pipeline
-pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
-
-# Your other functions and main code here...
+# Retry mechanism for loading the model
+def retry_load_model():
+    try:
+        return StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
+    except EnvironmentError 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
@@ -32,4 +38,3 @@ def main():
 if __name__ == '__main__':
     main()
 
-