reformat

app.py

@@ -152,58 +152,37 @@ def combine_components_slice(model, gd, im1, im2, indices=None, sample_method='d
 
 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)
+    result = gen_image_and_components(MODELS[model], GD[sample_method], im, sample_method=sample_method, num_images=1, device=device)
     return result.permute(1, 2, 0).numpy()
 
 
 def combine_images_demo(im1, im2, model):
     sample_method = 'ddim'
-    result = combine_components_slice(MODELS[model], GD[sample_method], im1, im2, indices='1,0,1,0', sample_method=sample_method, num_images=1)
+    result = combine_components_slice(MODELS[model], GD[sample_method], im1, im2, indices='1,0,1,0', sample_method=sample_method, num_images=1, device=device)
     return result.permute(1, 2, 0).numpy()
 
 
+def load_model(dataset, extra_kwargs={}, device='cuda'):
+    ckpt_path = download_model(dataset) 
 
+    model_kwargs = unet_model_defaults()
+    # model parameters
+    model_kwargs.update(extra_kwargs)
+    model = create_diffusion_model(**model_kwargs)
+    model.eval()
+    model.to(device)
 
-ckpt_path = download_model('clevr') # 'clevr_model.pt'
+    print(f'loading from {ckpt_path}')
+    checkpoint = th.load(ckpt_path, map_location='cpu')
 
-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()
+    model.load_state_dict(checkpoint)
+    return model
 
-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()
-# model parameters
-model_kwargs.update(dict(
-    enc_channels=128
-))
-celeb_model = create_diffusion_model(**model_kwargs)
-celeb_model.eval()
 
 device = 'cuda' if th.cuda.is_available() else 'cpu'
-celeb_model.to(device)
-
-print(f'loading from {ckpt_path}')
-checkpoint = th.load(ckpt_path, map_location='cpu')
-
-celeb_model.load_state_dict(checkpoint)
-
 
+clevr_model = load_model('clevr', extra_kwargs=dict(embed_dim=64, enc_channels=128), device=device)
+celeb_model = load_model('celebahq', extra_kwargs=dict(enc_channels=128), device=device)
 
 MODELS = {
     'CLEVR': clevr_model,
@@ -222,7 +201,7 @@ with gr.Blocks() as demo:
     
     gr.Markdown(
             """<h4>Decomposition and reconstruction of images</h4>""")
-    with gr.Row().style(equal_height=True):
+    with gr.Row():
         with gr.Column():
             with gr.Row():
                 decomp_input = gr.Image(type='numpy', label='Input')
@@ -230,19 +209,21 @@ with gr.Blocks() as demo:
                 decomp_model = gr.Radio(
                     ['CLEVR', 'CelebA-HQ'], type="value", label='Model',
                     value='CLEVR')
+                
+            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]
+                )
             
         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(
@@ -260,20 +241,21 @@ with gr.Blocks() as demo:
                 comb_model = gr.Radio(
                     ['CLEVR', 'CelebA-HQ'], type="value", label='Model',
                     value='CLEVR')
-                
+            
+            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]
+                )
             
 
         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],