Update app.py

app.py

@@ -35,20 +35,23 @@ except Exception as e:
 def generate_image_from_model(prompt):
     try:
         response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
+        response.raise_for_status()  # Raise an error for bad responses
         image_bytes = response.content
+        if not image_bytes:
+            raise ValueError("Empty image content received from the API")
         image = Image.open(io.BytesIO(image_bytes))
         return image
     except Exception as e:
-        st.error(f"Error occurs:{e}")
+        st.error(f"Error generating image from model: {e}")
         
+        return None
+
 def generate_image(prompt):
-    try:
-        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
-        return image
-    except Exception as e:
-        st.error(f"Error occurs:{e}")
+    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
+    image_bytes = response.content
+    image = Image.open(io.BytesIO(image_bytes))
+    return image
+
 # Streamlit app
 def main():
     try:
@@ -123,15 +126,18 @@ def main():
                         st.session_state.chat_history.append({"role": "bot", "content": generated_image})
 
                         # Display the generated image
-                        for index, chat in enumerate(st.session_state.chat_history):
-                            with st.chat_message(chat["role"]):
-                                if chat["role"] == "user":
-                                    st.markdown(user_input)
-                                elif chat["role"] == "bot":
-                                    st.image(generated_image, width=400)
+                        if generated_image is not None:
+                            for index, chat in enumerate(st.session_state.chat_history):
+                                with st.chat_message(chat["role"]):
+                                    if chat["role"] == "user":
+                                        st.markdown(user_input)
+                                    elif chat["role"] == "bot":
+                                        st.image(generated_image, width=400)
+                        else:
+                            st.error("Failed to generate image. Check logs for details.")
 
                     else:
-                        GOOGLE_API_KEY = "your_Gemini_Api_key"
+                        GOOGLE_API_KEY = "your_gemini_Api_key"
                         genai.configure(api_key=GOOGLE_API_KEY)
                         model = genai.GenerativeModel('gemini-1.0-pro')
                         prompt = user_input
@@ -160,12 +166,15 @@ def main():
             elif selected_model == "stabilityai/stable-diffusion-xl-base-1.0":
                 prompt = user_input
                 generated_image = generate_image_from_model(prompt)
-                for index, chat in enumerate(st.session_state.chat_history):
-                    with st.chat_message(chat["role"]):
-                        if chat["role"] == "user":
-                            st.markdown(user_input)
-                        elif chat["role"] == "bot":
-                            st.image(generated_image, width=400)
+                if generated_image is not None:
+                    for index, chat in enumerate(st.session_state.chat_history):
+                        with st.chat_message(chat["role"]):
+                            if chat["role"] == "user":
+                                st.markdown(user_input)
+                            elif chat["role"] == "bot":
+                                st.image(generated_image, width=400)
+                else:
+                    st.error("Failed to generate image. Check logs for details.")
 
             else:
                 try: