Update app.py

app.py

@@ -10,8 +10,8 @@ from torchvision.transforms import functional as F
 from diffusers import (
     AutoPipelineForInpainting,
 )
-from generate_dataset import outpainting_generator_rectangle
-
+from generate_dataset import outpainting_generator_rectangle, merge_images_horizontally
+from ddim_with_prob import DDIMSchedulerCustom
 
 transform = transforms.Compose([
         transforms.ToPILImage(),
@@ -23,15 +23,23 @@ def pref_inpainting(image,
                     mask_random_start,
                     steps,
                     ):
-    with open("/data0/kendong/Diffusions/zero123-live/configs/imagereward_train_configs.yaml") as file:
+    with open("./configs/paintreward_train_configs.yaml") as file:
         config_dict= yaml.safe_load(file)
         config = munchify(config_dict)
 
-    pipe = AutoPipelineForInpainting.from_pretrained(
-                '/data1/kendong/joint-rl-diffusion/alignment_log/exp_reward_group_regression_all_1w_1.6boundary/iteration_2560', num_inference_steps=steps)
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    pipe = pipe.to(device)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    
+    pipe_ours = AutoPipelineForInpainting.from_pretrained(
+                './model_ckpt', torch_dtype=torch.float16, variant='fp16')
+    pipe_ours.scheduler = DDIMSchedulerCustom.from_config(pipe_ours.scheduler.config)
 
+    pipe_runway = AutoPipelineForInpainting.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant='fp16')
+    
+    
+    pipe_ours = pipe_ours.to(device)
+    pipe_runway = pipe_runway.to(device)
+    print('Loading pipeline')
    
     color, mask = outpainting_generator_rectangle(image, box_width_ratio/100, mask_random_start)
     mask = mask.convert('L')
@@ -47,12 +55,16 @@ def pref_inpainting(image,
 
 
     color, mask = transform(color), transform(mask)
-    res = pipe(prompt='', image=color, mask_image=mask, eta=config.eta).images[0]
-    
+    res_ours = pipe_ours(prompt='', image=color, mask_image=mask, eta=config.eta).images[0]
+    print('Running inference ours')
+    res_runway = pipe_runway(prompt="", image=color, mask_image=mask).images[0]
+    print('Running inference runway')
     
     # res.save(os.path.join('./',  'test.png'))
-    
-    return color, res
+    res_ours = merge_images_horizontally(color, res_ours)
+    res_runway = merge_images_horizontally(color, res_runway)
+
+    return res_ours, res_runway
 
 
 inputs = [
@@ -63,29 +75,31 @@ inputs = [
 ]
 
 outputs = [
-    gr.Image(type="pil", image_mode="RGBA", label='Mask RGB Image', container=True, width="65%"),
-    gr.Image(type="pil", image_mode="RGBA", label='Results', container=True, width="65%"), 
+    gr.Image(type="pil", image_mode="RGBA", label='PrefPaint', container=True, width="100%"),
+    gr.Image(type="pil", image_mode="RGBA", label='RunwayPaint', container=True, width="100%"), 
 ]
 
+files = os.listdir("./assets")
 examples = [
-    ["/data0/kendong/Diffusions/zero123-live/test_demo/assets/ILSVRC2012_test_00000003.JPEG", 35, 125, 50],
-    ["/data0/kendong/Diffusions/zero123-live/test_demo/assets/ILSVRC2012_test_00000181.JPEG", 35, 125, 50],
-    ["/data0/kendong/Diffusions/zero123-live/test_demo/assets/ILSVRC2012_test_00002334.JPEG", 35, 125, 50],
-    ["/data0/kendong/Diffusions/zero123-live/test_demo/assets/ILSVRC2012_test_00002613.JPEG", 35, 125, 50],
+    [f"./assets/{file_name}", 35, 125, 50] for file_name in files
 ]
 
 
-        
-iface = gr.Interface(
-    fn=pref_inpainting,
-    inputs=inputs,
-    outputs=outputs,
-    title="Inpainting with Human Preference (Utilizing Free CPU Resources)",
-    description="Upload an image and start your inpainting (currently only supporting outpainting masks; other mask types coming soon).",
-    theme="default",
-    examples= examples,
-    allow_flagging="never"
-)
+with gr.Blocks() as demo:
+   
+
+    iface = gr.Interface(
+        fn=pref_inpainting,
+        inputs=inputs,
+        outputs=outputs,
+        title="Inpainting with Human Preference (Utilizing Free CPU Resources)",
+        description="Upload an image and start your inpainting (currently only supporting outpainting masks; other mask types coming soon).",
+        theme="default",
+        examples=examples,
+        # allow_flagging="never"
+    )
+   
 
+    # iface.launch()
 
-iface.launch(share=True)
+demo.launch()