updated for both decomp and combine

app.py

@@ -98,46 +98,6 @@ gd = SpacedDiffusion(spaced_ts, rescale_timesteps=True, original_num_steps=num_t
 GD['ddim'] = gd
 
 
-
-# ckpt_path = download_model('clevr') # 'clevr_model.pt'
-
-# model_kwargs = unet_model_defaults()
-# # model parameters
-# model_kwargs.update(dict(
-#     emb_dim=64,
-#     enc_channels=128
-# ))
-# clevr_model = create_diffusion_model(**model_kwargs)
-# clevr_model.eval()
-
-# device = 'cuda' if th.cuda.is_available() else 'cpu'
-# clevr_model.to(device)
-
-# print(f'loading from {ckpt_path}')
-# checkpoint = th.load(ckpt_path, map_location='cpu')
-
-# clevr_model.load_state_dict(checkpoint)
-
-
-
-# img_input =  gr.inputs.Image(type="numpy", label="Input")
-# img_output = gr.outputs.Image(type="numpy", label="Output")
-
-# gr.Interface(
-#     decompose_image,
-#     inputs=img_input,
-#     outputs=img_output,
-#     examples=[
-#         "sample_images/clevr_im_10.png",
-#         "sample_images/clevr_im_25.png",
-#     ],
-
-# ).launch()
-
-
-
-
-
 def combine_components_slice(model, gd, im1, im2, indices=None, sample_method='ddim', device='cuda', num_images=4, model_kwargs={}, desc='', save_dir='', dataset='clevr', image_size=64):
     """Combine by adding components together
     """
@@ -188,13 +148,43 @@ def combine_components_slice(model, gd, im1, im2, indices=None, sample_method='d
 
     return sample[0].cpu()
 
-def combine_images(im1, im2):
+
+
+def decompose_image_demo(im, model):
+    sample_method = 'ddim'
+    result = gen_image_and_components(MODELS[model], GD[sample_method], im, sample_method=sample_method, num_images=1)
+    return result.permute(1, 2, 0).numpy()
+
+
+def combine_images_demo(im1, im2, model):
     sample_method = 'ddim'
-    result = combine_components_slice(celeb_model, GD[sample_method], im1, im2, indices='1,0,1,0', sample_method=sample_method, num_images=1, device=device)
+    result = combine_components_slice(MODELS[model], GD[sample_method], im1, im2, indices='1,0,1,0', sample_method=sample_method, num_images=1)
     return result.permute(1, 2, 0).numpy()
 
 
 
+
+ckpt_path = download_model('clevr') # 'clevr_model.pt'
+
+model_kwargs = unet_model_defaults()
+# model parameters
+model_kwargs.update(dict(
+    emb_dim=64,
+    enc_channels=128
+))
+clevr_model = create_diffusion_model(**model_kwargs)
+clevr_model.eval()
+
+device = 'cuda' if th.cuda.is_available() else 'cpu'
+clevr_model.to(device)
+
+print(f'loading from {ckpt_path}')
+checkpoint = th.load(ckpt_path, map_location='cpu')
+
+clevr_model.load_state_dict(checkpoint)
+
+
+
 ckpt_path = download_model('celebahq') # 'celeb_model.pt'
 
 model_kwargs = unet_model_defaults()
@@ -213,21 +203,84 @@ checkpoint = th.load(ckpt_path, map_location='cpu')
 
 celeb_model.load_state_dict(checkpoint)
 
-# Recombination
-
 
-img_input =  gr.inputs.Image(type="numpy", label="Input")
-img_input2 =  gr.inputs.Image(type="numpy", label="Input")
 
-img_output = gr.outputs.Image(type="numpy", label="Output")
+MODELS = {
+    'CLEVR': clevr_model,
+    'CelebA-HQ': celeb_model
+}
 
-gr.Interface(
-    combine_images,
-    inputs=[img_input, img_input2],
-    outputs=img_output,
-    examples=[
-        ["sample_images/celebahq_im_15.jpg",
-        "sample_images/celebahq_im_21.jpg"]
-    ]
-).launch()
 
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """<h1 style="text-align: center;"><b>Unsupervised Compositional Image Decomposition with Diffusion Models
+</b> - <a href="https://jsu27.github.io/decomp-diffusion-web/">Project Page</a></h1>""")
+    
+    gr.Markdown(
+        """<p style="font-size: 18px;">We introduce Decomp Diffusion, an unsupervised approach that discovers compositional concepts from images, represented by diffusion models.
+</p>""")
+    
+    gr.Markdown(
+            """<h4>Decomposition and reconstruction of images</h4>""")
+    with gr.Row().style(equal_height=True):
+        with gr.Column():
+            with gr.Row():
+                decomp_input = gr.Image(type='numpy', label='Input')
+            with gr.Row():
+                decomp_model = gr.Radio(
+                    ['CLEVR', 'CelebA-HQ'], type="value", label='Model',
+                    value='CLEVR')
+            
+        with gr.Column():
+            decomp_output = gr.Image(type='numpy')
+            decomp_button = gr.Button("Generate")
+    with gr.Row():
+
+        # image_examples = [os.path.join(os.path.dirname(__file__), 'sample_images/clevr_im_10.png'), 'CLEVR']
+        decomp_examples = [['sample_images/clevr_im_10.png', 'CLEVR'],
+                            ['sample_images/celebahq_im_15.jpg', 'CelebA-HQ']]
+        decomp_img_examples = gr.Examples(
+            examples=decomp_examples,
+            inputs=[decomp_input, decomp_model]
+        )
+
+
+    gr.Markdown(
+            """<h4>Combination of images</h4>""")
+    with gr.Row().style(equal_height=True):
+        with gr.Column(scale=2):
+
+            with gr.Row():
+                with gr.Column():
+                    comb_input1 = gr.Image(type='numpy', label='Input 1')
+                with gr.Column():
+                    comb_input2 = gr.Image(type='numpy', label='Input 2')
+
+            with gr.Row():
+                comb_model = gr.Radio(
+                    ['CLEVR', 'CelebA-HQ'], type="value", label='Model',
+                    value='CLEVR')
+                
+            
+
+        with gr.Column(scale=1):
+            comb_output = gr.Image(type='numpy')
+            comb_button = gr.Button("Generate")
+    with gr.Row():
+
+        comb_examples = [['sample_images/clevr_im_10.png', 'sample_images/clevr_im_25.png', 'CLEVR'],
+                            ['sample_images/celebahq_im_15.jpg', 'sample_images/celebahq_im_21.jpg', 'CelebA-HQ']]
+        comb_img_examples = gr.Examples(
+            examples=comb_examples,
+            inputs=[comb_input1, comb_input2, comb_model]
+        )
+
+    decomp_button.click(decompose_image_demo,
+                       inputs=[decomp_input, decomp_model],
+                       outputs=decomp_output)
+    comb_button.click(combine_images_demo,
+                       inputs=[comb_input1, comb_input2, comb_model],
+                       outputs=comb_output)
+
+
+demo.launch()