Update generator_function/image_function.py

generator_function/image_function.py

@@ -283,13 +283,13 @@ def handle_image_generation_optimized(
 
                         image_bytes_list.append((f"image_{i + 1}.png", b))
 
-                        st.image(b, use_container_width=True)
+                        st.image(b, width='stretch')
                         st.download_button(
                             f"Download image ",
                             b,
                             file_name=f"image_{i + 1}.png",
                             mime="image/png",
-                            use_container_width=True,
+                            width='stretch',
                         )
                     except Exception as e:
                         st.error(f"Display failed: {e}")
@@ -307,7 +307,7 @@ def handle_image_generation_optimized(
                     data=zip_buffer,
                     file_name="all_images.zip",
                     mime="application/zip",
-                    use_container_width=True,
+                    width='stretch',
                 )
 
         else: