add

app.py

@@ -22,7 +22,7 @@ model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir
 # sys.path.insert(0, model_folder)
 # print(sys.path)
 
-from ultrafusion_utils import load_model, run_ultrafusion
+from ultrafusion_utils import load_model, run_ultrafusion, check_input
 
 to_tensor = ToTensor()
 to_pil = ToPILImage()
@@ -38,36 +38,23 @@ else:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-@spaces.GPU(duration=10) #[uncomment to use ZeroGPU]
+@spaces.GPU(duration=60) #[uncomment to use ZeroGPU]
 def infer(
     under_expo_img,
     over_expo_img,
-    # progress=gr.Progress(track_tqdm=True),
+    num_inference_steps
 ):
-    # if randomize_seed:
-    #     seed = random.randint(0, MAX_SEED)
-    # generator = torch.Generator().manual_seed(seed)
-
-    # image = pipe(
-    #     prompt=prompt,
-    #     negative_prompt=negative_prompt,
-    #     guidance_scale=guidance_scale,
-    #     num_inference_steps=num_inference_steps,
-    #     width=width,
-    #     height=height,
-    #     generator=generator,
-    # ).images[0]
-
     print(under_expo_img.size)
     print("reciving image")
 
-    under_expo_img = under_expo_img.resize([1500, 1000])
-    over_expo_img = over_expo_img.resize([1500, 1000])
+    # under_expo_img = under_expo_img.resize([1500, 1000])
+    # over_expo_img = over_expo_img.resize([1500, 1000])
+    under_expo_img, over_expo_img = check_input(under_expo_img, over_expo_img, max_l=1500)
 
     ue = to_tensor(under_expo_img).unsqueeze(dim=0).to("cuda")
     oe = to_tensor(over_expo_img).unsqueeze(dim=0).to("cuda")
 
-    out = run_ultrafusion(ue, oe, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, consistent_start=None)
+    out = run_ultrafusion(ue, oe, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, steps=num_inference_steps, consistent_start=None)
 
     out = out.clamp(0, 1).squeeze()
     out_pil = to_pil(out)
@@ -143,84 +130,26 @@ with gr.Blocks(css=css) as demo:
                           width=IMG_W*2,
                           height=IMG_H*2,
                         )
-
-        # with gr.Row():
-        # prompt = gr.Text(
-        #     label="Prompt",
-        #     show_label=False,
-        #     max_lines=1,
-        #     placeholder="Enter your prompt",
-        #     container=False,
-        # )
-        # negative_prompt = gr.Text(
-        #     label="Negative prompt",
-        #     max_lines=1,
-        #     placeholder="Enter a negative prompt",
-        #     visible=False,
-        # )
-        # with gr.Accordion("Advanced Settings", open=False):
-        #     negative_prompt = gr.Text(
-        #         label="Negative prompt",
-        #         max_lines=1,
-        #         placeholder="Enter a negative prompt",
-        #         visible=False,
-        #     )
-
-        #     seed = gr.Slider(
-        #         label="Seed",
-        #         minimum=0,
-        #         maximum=MAX_SEED,
-        #         step=1,
-        #         value=0,
-        #     )
-
-        #     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-        #     with gr.Row():
-        #         width = gr.Slider(
-        #             label="Width",
-        #             minimum=256,
-        #             maximum=MAX_IMAGE_SIZE,
-        #             step=32,
-        #             value=1024,  # Replace with defaults that work for your model
-        #         )
-
-        #         height = gr.Slider(
-        #             label="Height",
-        #             minimum=256,
-        #             maximum=MAX_IMAGE_SIZE,
-        #             step=32,
-        #             value=1024,  # Replace with defaults that work for your model
-        #         )
-
-        #     with gr.Row():
-        #         guidance_scale = gr.Slider(
-        #             label="Guidance scale",
-        #             minimum=0.0,
-        #             maximum=10.0,
-        #             step=0.1,
-        #             value=0.0,  # Replace with defaults that work for your model
-        #         )
-
-        #         num_inference_steps = gr.Slider(
-        #             label="Number of inference steps",
-        #             minimum=1,
-        #             maximum=50,
-        #             step=1,
-        #             value=2,  # Replace with defaults that work for your model
-        #         )
+        with gr.Accordion("Advanced Settings", open=True):
+            num_inference_steps = gr.Slider(
+                label="Number of inference steps",
+                minimum=2,
+                maximum=50,
+                step=1,
+                value=20,  # Replace with defaults that work for your model
+                interactive=True
+            )
        
         gr.Examples(
             examples=examples, 
             inputs=[under_expo_img, over_expo_img], 
             label="Examples",
             # examples_per_page=10,
-            cache_examples=False,
-            # fn=infer,
+            cache_examples=True,
             )
         # gr.Markdown(_CITE_)
     run_button.click(fn=infer,
-                     inputs=[under_expo_img, over_expo_img],
+                     inputs=[under_expo_img, over_expo_img, num_inference_steps],
                      outputs=[result,],
                     )