diff --git a/package/boot/uboot-mediatek/patches/200-cmd-add-imsz-and-imszb.patch b/package/boot/uboot-mediatek/patches/200-cmd-add-imsz-and-imszb.patch
index 27cea2fa4e3..9872ccbf268 100644
--- a/package/boot/uboot-mediatek/patches/200-cmd-add-imsz-and-imszb.patch
+++ b/package/boot/uboot-mediatek/patches/200-cmd-add-imsz-and-imszb.patch
@@ -1,17 +1,16 @@
 --- a/cmd/bootm.c
 +++ b/cmd/bootm.c
-@@ -260,6 +260,67 @@ U_BOOT_CMD(
+@@ -260,6 +260,76 @@ U_BOOT_CMD(
  /* iminfo - print header info for a requested image */
  /*******************************************************************/
  #if defined(CONFIG_CMD_IMI)
-+#if defined(CONFIG_FIT)
 +#define SECTOR_SHIFT 9
 +static int image_totalsize(struct cmd_tbl *cmdtp, int flag, int argc,
 +			   char *const argv[], short int in_blocks)
 +{
 +	ulong addr;
-+	void *fit;
-+	int bsize, tsize;
++	void *hdr;
++	uint32_t bsize, tsize = 0;
 +	char buf[16];
 +
 +	if (argc >= 2)
@@ -19,9 +18,20 @@
 +	else
 +		addr = image_load_addr;
 +
-+	fit = (void *)map_sysmem(addr, 0);
-+	tsize = fit_get_totalsize(fit);
-+	unmap_sysmem(fit);
++	hdr = (void *)map_sysmem(addr, 0);
++
++	switch (genimg_get_format(hdr)) {
++	case IMAGE_FORMAT_LEGACY:
++		if(CONFIG_IS_ENABLED(LEGACY_IMAGE_FORMAT))
++			tsize = image_get_image_size(hdr);
++		break;
++	case IMAGE_FORMAT_FIT:
++		if(CONFIG_IS_ENABLED(FIT))
++			tsize = fit_get_totalsize(hdr);
++		break;
++	}
++
++	unmap_sysmem(hdr);
 +	if (tsize == 0)
 +		return 1;
 +
@@ -64,7 +74,6 @@
 +	"addr [maxhdrlen] [varname]\n"
 +);
 +
-+#endif
  static int do_iminfo(struct cmd_tbl *cmdtp, int flag, int argc,
  		     char *const argv[])
  {