fix: Update process_sr function to return Image instead of List[np.ndarray] and adjust output gallery type

app.py

@@ -322,7 +322,7 @@ def process_sr(
     cfg_scale: float,
     seed: int,
     model_choice: str,
-    ) -> List[np.ndarray]:
+    ) -> Image.Image:
     process_size = 512
     resize_preproc = transforms.Compose([
         transforms.Resize(process_size, interpolation=transforms.InterpolationMode.BILINEAR),
@@ -349,8 +349,6 @@ def process_sr(
     input_image = input_image.resize((input_image.size[0] // 8 * 8, input_image.size[1] // 8 * 8))
     width, height = input_image.size
     resize_flag = True  #
-
-    images = []
     
     # Choose pipeline based on model selection - prioritize dit4sr_f
     if model_choice == "dit4sr_q" and pipeline_dit4sr_q is not None:
@@ -380,8 +378,8 @@ def process_sr(
     except Exception as e:
         print(f"Error during inference: {e}")
         image = Image.new(mode="RGB", size=(512, 512))
-    images.append(np.array(image))
-    return images
+    
+    return image
 
 
 
@@ -451,7 +449,7 @@ with block:
                 scale_factor = gr.Number(label="SR Scale", value=4)
             gr.Examples(examples=exaple_images, inputs=[input_image])
         with gr.Column():
-            result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery", columns=1)
+            result_gallery = gr.Image(label="Output", show_label=False, elem_id="gallery", type="pil", format="png")
             with gr.Row():
                 run_llava_button = gr.Button(value="Run LLAVA")
                 run_sr_button = gr.Button(value="Run DiT4SR")