diff --git a/main.py b/main.py
index 908ff7af..9f1c8928 100644
--- a/main.py
+++ b/main.py
@@ -41,10 +41,10 @@ async def run(server, address='', port=8188, verbose=True, call_on_start=None):
     await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop())
 
 def hijack_progress(server):
-    def hook(value, total, preview_bytes_jpeg):
+    def hook(value, total, preview_image_bytes):
         server.send_sync("progress", { "value": value, "max": total}, server.client_id)
-        if preview_bytes_jpeg is not None:
-            server.send_sync(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes_jpeg, server.client_id)
+        if preview_image_bytes is not None:
+            server.send_sync(BinaryEventTypes.PREVIEW_IMAGE, preview_image_bytes, server.client_id)
             pass
     comfy.utils.set_progress_bar_global_hook(hook)