Refactor the UI layout of ImagenScreen to solve the navigation flicker

firebaseai/ImagenScreen/ImagenScreen.swift

@@ -34,43 +34,42 @@ struct ImagenScreen: View {
 
   var body: some View {
     ZStack {
-      VStack {
-        InputField("Enter a prompt to generate an image", text: $viewModel.userInput) {
-          Image(
-            systemName: viewModel.inProgress ? "stop.circle.fill" : "paperplane.circle.fill"
-          )
-          .font(.title)
-        }
-        .focused($focusedField, equals: .message)
-        .onSubmit { sendOrStop() }
+      ScrollView {
+        VStack {
+          InputField("Enter a prompt to generate an image", text: $viewModel.userInput) {
+            Image(
+              systemName: viewModel.inProgress ? "stop.circle.fill" : "paperplane.circle.fill"
+            )
+            .font(.title)
+          }
+          .focused($focusedField, equals: .message)
+          .onSubmit { sendOrStop() }
 
-        ScrollView {
           let spacing: CGFloat = 10
           LazyVGrid(columns: [
-            GridItem(.fixed(UIScreen.main.bounds.width / 2 - spacing), spacing: spacing),
-            GridItem(.fixed(UIScreen.main.bounds.width / 2 - spacing), spacing: spacing),
+            GridItem(.flexible(), spacing: spacing),
+            GridItem(.flexible(), spacing: spacing),
           ], spacing: spacing) {
             ForEach(viewModel.images, id: \.self) { image in
               Image(uiImage: image)
                 .resizable()
-                .aspectRatio(contentMode: .fill)
-                .frame(width: UIScreen.main.bounds.width / 2 - spacing,
-                       height: UIScreen.main.bounds.width / 2 - spacing)
+                .aspectRatio(1, contentMode: .fill)
                 .cornerRadius(12)
                 .clipped()
             }
           }
           .padding(.horizontal, spacing)
         }
       }
+      .navigationTitle("Imagen example")
+      .onAppear {
+        focusedField = .message
+      }
+
       if viewModel.inProgress {
         ProgressOverlay()
       }
     }
-    .navigationTitle("Imagen example")
-    .onAppear {
-      focusedField = .message
-    }
   }
 
   private func sendMessage() {
@@ -92,18 +91,23 @@ struct ImagenScreen: View {
 struct ProgressOverlay: View {
   var body: some View {
     ZStack {
-      RoundedRectangle(cornerRadius: 16)
-        .fill(Material.ultraThinMaterial)
-        .frame(width: 120, height: 100)
-        .shadow(radius: 8)
+      Color.black.opacity(0.3)
+        .ignoresSafeArea()
+
+      ZStack {
+        RoundedRectangle(cornerRadius: 16)
+          .fill(Material.ultraThinMaterial)
+          .frame(width: 120, height: 100)
+          .shadow(radius: 8)
 
-      VStack(spacing: 12) {
-        ProgressView()
-          .scaleEffect(1.5)
+        VStack(spacing: 12) {
+          ProgressView()
+            .scaleEffect(1.5)
 
-        Text("Loading...")
-          .font(.subheadline)
-          .foregroundColor(.secondary)
+          Text("Loading...")
+            .font(.subheadline)
+            .foregroundColor(.secondary)
+        }
       }
     }
   }