diff --git a/javascript/ui.js b/javascript/ui.js
index 4f89a66d4..d48e13ef4 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -317,9 +317,9 @@ function selectCheckpoint(name) {
     gradioApp().getElementById('change_checkpoint').click();
 }
 
-function currentImg2imgSourceResolution(w, h, scaleBy) {
-    var img = gradioApp().querySelector('#mode_img2img > div[style="display: block;"] img');
-    return img ? [img.naturalWidth, img.naturalHeight, scaleBy] : [0, 0, scaleBy];
+function currentImg2imgSourceResolution(w, h) {
+    var img = gradioApp().querySelector('#mode_img2img > div[style="display: block;"] :is(img, canvas)');
+    return img ? [img.naturalWidth || img.width, img.naturalHeight || img.height] : [0, 0];
 }
 
 function updateImg2imgResizeToTextAfterChangingImage() {
diff --git a/modules/ui.py b/modules/ui.py
index 442d614e4..5be344d22 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -774,9 +774,9 @@ def create_ui():
             res_switch_btn.click(fn=None, _js="function(){switchWidthHeight('img2img')}", inputs=None, outputs=None, show_progress=False)
 
             detect_image_size_btn.click(
-                fn=lambda w, h, _: (w or gr.update(), h or gr.update()),
+                fn=lambda w, h: (w or gr.update(), h or gr.update()),
                 _js="currentImg2imgSourceResolution",
-                inputs=[dummy_component, dummy_component, dummy_component],
+                inputs=[dummy_component, dummy_component],
                 outputs=[width, height],
                 show_progress=False,
             )