Update pages/Visualization.py

pages/Visualization.py

@@ -23,19 +23,20 @@ image_idea = st.text_input('Предложите свою тему для ген
 image_gen__btn = st.button('Генерировать изображение')
 if image_gen__btn:
     with st.spinner('...'):
-        image_bytes = generate_img({"inputs": image_idea}) 
-        image_raw = io.BytesIO(image_bytes)        
-        st.success('Готово')        
-    if image_raw is not None:        
-        st.image(image_raw, width=600)     
-        st.markdown('## Опишите фотографию на английском языке') 
-        st.markdown('## План ответа поможет вам:') 
-        st.markdown('+ the place;') 
-        st.markdown('+ the action;') 
-        st.markdown('+ the person’s appearance;') 
-        st.markdown('+ whether you like the picture or not;') 
-        st.markdown('+ why.') 
-        st.markdown('Start with: “I’d like to describe this picture. The picture shows …” ')
+        try:
+            image_bytes = generate_img({"inputs": image_idea}) 
+            image_raw = io.BytesIO(image_bytes) 
+        #except UnidentifiedImageError:            
+            st.success('Готово')            
+            st.image(image_raw)     
+            st.markdown('## Опишите фотографию на английском языке') 
+            st.markdown('## План ответа поможет вам:') 
+            st.markdown('+ the place;') 
+            st.markdown('+ the action;') 
+            st.markdown('+ the person’s appearance;') 
+            st.markdown('+ whether you like the picture or not;') 
+            st.markdown('+ why.') 
+            st.markdown('Start with: “I’d like to describe this picture. The picture shows …” ')
 
 st.divider()