Update app.py

app.py

@@ -5,13 +5,31 @@ from PIL import Image
 import torch
 
 # Initialize the text generation model
-text_generator = pipeline('text-generation', model='gpt2')
+def initialize_text_generator():
+    try:
+        text_generator = pipeline('text-generation', model='gpt2')
+    except Exception as e:
+        st.error(f"Error loading text generation model: {e}")
+        return None
+    return text_generator
 
 # Initialize the image generation model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-image_generator = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
+def initialize_image_generator():
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        image_generator = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
+    except Exception as e:
+        st.error(f"Error loading image generation model: {e}")
+        return None
+    return image_generator
+
+text_generator = initialize_text_generator()
+image_generator = initialize_image_generator()
 
 def generate_blog(title):
+    if text_generator is None or image_generator is None:
+        return "Failed to load models", None
+    
     # Generate blog content
     blog_content = text_generator(title, max_length=500, num_return_sequences=1)[0]['generated_text']
     
@@ -28,11 +46,11 @@ title = st.text_input('Enter the title of your blog:')
 if title:
     with st.spinner('Generating blog content and image...'):
         blog_content, image = generate_blog(title)
-        st.success('Blog generated successfully!')
-        st.subheader('Blog Content')
-        st.write(blog_content)
-        st.subheader('Generated Image')
-        st.image(image, caption='Generated Image')
-
-if __name__ == '__main__':
-    st.run()
+        if blog_content == "Failed to load models":
+            st.error(blog_content)
+        else:
+            st.success('Blog generated successfully!')
+            st.subheader('Blog Content')
+            st.write(blog_content)
+            st.subheader('Generated Image')
+            st.image(image, caption='Generated Image')