celeb combination demo

app.py

@@ -98,40 +98,136 @@ gd = SpacedDiffusion(spaced_ts, rescale_timesteps=True, original_num_steps=num_t
 GD['ddim'] = gd
 
 
-# !wget https://www.dropbox.com/s/bqpc3ymstz9m05z/clevr_model.pt
-# load model
 
-ckpt_path = download_model('clevr') # 'clevr_model.pt'
+# 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
+    """
+    assert sample_method in ('ddpm', 'ddim')
+    
+    im1 = get_pil_im(im1, resolution=image_size).to(device)
+    im2 = get_pil_im(im2, resolution=image_size).to(device)
+
+    latent1 = model.encode_latent(im1)
+    latent2 = model.encode_latent(im2)
+
+    num_comps = model.num_components 
+
+    # get latent slices
+    if indices == None:
+        half = num_comps // 2
+        indices = [1] * half + [0] * half # first half 1, second half 0
+        indices = th.Tensor(indices) == 1
+        indices = indices.reshape(num_comps, 1)
+    elif type(indices) == str:
+        indices = indices.split(',')
+        indices = [int(ind) for ind in indices]
+        indices = th.Tensor(indices).reshape(-1, 1) == 1
+    assert len(indices) == num_comps
+    indices = indices.to(device)
+
+    latent1 = latent1.reshape(num_comps, -1).to(device)
+    latent2 = latent2.reshape(num_comps, -1).to(device)
+
+    combined_latent = th.where(indices, latent1, latent2)
+    combined_latent = combined_latent.reshape(1, -1)
+    model_kwargs['latent'] = combined_latent
+    
+    sample_loop_func = gd.p_sample_loop if sample_method == 'ddpm' else gd.ddim_sample_loop
+    if sample_method == 'ddim':
+        model = gd._wrap_model(model)
+
+    # sampling loop
+    sample = sample_loop_func(
+            model,
+            (1, 3, image_size, image_size),
+            device=device,
+            clip_denoised=True,
+            progress=True,
+            model_kwargs=model_kwargs,
+            cond_fn=None,
+        )[:1]
+
+    return sample[0].cpu()
+
+def combine_images(im1, im2):
+    sample_method = 'ddim'
+    result = combine_components_slice(clevr_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('celebahq') # 'celeb_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()
+celeb_model = create_diffusion_model(**model_kwargs)
+celeb_model.eval()
 
 device = 'cuda' if th.cuda.is_available() else 'cpu'
-clevr_model.to(device)
+celeb_model.to(device)
 
 print(f'loading from {ckpt_path}')
 checkpoint = th.load(ckpt_path, map_location='cpu')
 
-clevr_model.load_state_dict(checkpoint)
+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")
 
 gr.Interface(
-    decompose_image,
-    inputs=img_input,
+    combine_images,
+    inputs=[img_input, img_input2],
     outputs=img_output,
     examples=[
-        os.path.join(os.path.dirname(__file__), "sample_images/clevr_im_10.png"),
-        os.path.join(os.path.dirname(__file__), "sample_images/clevr_im_25.png"),
-    ],
-
+        ["sample_images/celebahq_im_15.jpg",
+        "sample_images/celebahq_im_21.jpg"]
+    ]
 ).launch()
+

download.py

@@ -7,7 +7,7 @@ from tqdm.auto import tqdm
 
 MODEL_PATHS = {
     "clevr": "https://www.dropbox.com/s/bqpc3ymstz9m05z/clevr_model.pt",
-    "celebahq": ""
+    "celebahq": "https://www.dropbox.com/s/687wuamoud4cs9x/celeb_model.pt"
 }
 
 DATA_PATHS = {