Fix arguments

app.py

@@ -82,7 +82,12 @@ def check(input_image):
         raise gr.Error("Please provide an image to restore.")
 
 @spaces.GPU(duration=420)
-def stage1_process(input_image, gamma_correction):
+def stage1_process(
+    input_image,
+    gamma_correction,
+    diff_dtype,
+    ae_dtype
+):
     print('stage1_process ==>>')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
@@ -664,7 +669,9 @@ with gr.Blocks() as interface:
         input_image
     ], outputs = [], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
         input_image,
-        gamma_correction
+        gamma_correction,
+        diff_dtype,
+        ae_dtype
     ], outputs=[
         denoise_image,
         denoise_information