init

app.py

@@ -0,0 +1,137 @@
+import os
+import numpy as np
+import torch as th
+from imageio import imread
+from skimage.transform import resize as imresize
+
+from ema_pytorch import EMA
+from decomp_diffusion.model_and_diffusion_util import *
+from decomp_diffusion.diffusion.respace import SpacedDiffusion
+from decomp_diffusion.gen_image import *
+from download import download_model
+
+# fix randomness
+th.manual_seed(0)
+np.random.seed(0)
+
+import gradio as gr
+
+
+def get_pil_im(im, resolution=64):
+    im = imresize(im, (resolution, resolution))[:, :, :3]
+    im = th.Tensor(im).permute(2, 0, 1)[None, :, :, :].contiguous().cuda()
+    return im
+
+
+# generate image components and reconstruction
+def gen_image_and_components(model, gd, im, num_components=4, sample_method='ddim', batch_size=1, image_size=64, device='cuda', num_images=1):
+    """Generate row of orig image, individual components, and reconstructed image"""
+    orig_img = get_pil_im(im, resolution=image_size)
+    latent = model.encode_latent(orig_img)
+    model_kwargs = {'latent': latent}
+
+    assert sample_method in ('ddpm', 'ddim')
+    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)
+        
+    # generate imgs
+    for i in range(num_images):
+        all_samples = [orig_img]
+        # individual components
+        for j in range(num_components):
+            model_kwargs['latent_index'] = j
+            sample = sample_loop_func(
+                model,
+                (batch_size, 3, image_size, image_size),
+                device=device,
+                clip_denoised=True,
+                progress=True,
+                model_kwargs=model_kwargs,
+                cond_fn=None,
+            )[:batch_size]
+
+            # save indiv comp
+            all_samples.append(sample)
+        # reconstruction
+        model_kwargs['latent_index'] = None
+        sample = sample_loop_func(
+            model,
+            (batch_size, 3, image_size, image_size),
+            device=device,
+            clip_denoised=True,
+            progress=True,
+            model_kwargs=model_kwargs,
+            cond_fn=None,
+        )[:batch_size]
+        # save indiv reconstruction
+        all_samples.append(sample)
+
+        samples = th.cat(all_samples, dim=0).cpu()   
+        grid = make_grid(samples, nrow=samples.shape[0], padding=0)
+        return grid
+        
+
+def decompose_image(im):
+    sample_method = 'ddim'
+    result = gen_image_and_components(clevr_model, GD[sample_method], im, sample_method=sample_method, num_images=1)
+    return result.permute(1, 2, 0).numpy()
+
+
+# load diffusion
+GD = {} # diffusion objects for ddim and ddpm
+diffusion_kwargs = diffusion_defaults()
+gd = create_gaussian_diffusion(**diffusion_kwargs)
+GD['ddpm'] = gd
+
+# set up ddim sampling
+desired_timesteps = 50 
+num_timesteps = diffusion_kwargs['steps']
+
+spacing = num_timesteps // desired_timesteps
+spaced_ts = list(range(0, num_timesteps + 1, spacing))
+betas = get_named_beta_schedule(diffusion_kwargs['noise_schedule'], num_timesteps)
+diffusion_kwargs['betas'] = betas
+del diffusion_kwargs['steps'], diffusion_kwargs['noise_schedule']
+gd = SpacedDiffusion(spaced_ts, rescale_timesteps=True, original_num_steps=num_timesteps, **diffusion_kwargs)
+
+GD['ddim'] = gd
+
+
+# !wget https://www.dropbox.com/s/bqpc3ymstz9m05z/clevr_model.pt
+# load model
+
+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'
+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=[
+        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"),
+    ],
+
+).launch()

download.py

@@ -0,0 +1,44 @@
+
+import os
+from typing import Optional
+
+import requests
+from tqdm.auto import tqdm
+
+MODEL_PATHS = {
+    "clevr": "https://www.dropbox.com/s/bqpc3ymstz9m05z/clevr_model.pt",
+    "celebahq": ""
+}
+
+DATA_PATHS = {
+    "clevr": "",
+    "clevr_toy": ""
+}
+
+def download_model(
+    dataset: str,
+    cache_dir: Optional[str] = None,
+    chunk_size: int = 4096,
+) -> str:
+    if dataset not in MODEL_PATHS:
+        raise ValueError(
+            f"Unknown dataset name {dataset}. Known names are: {MODEL_PATHS.keys()}."
+        )
+    if cache_dir is None:
+        cache_dir = './'
+    url = MODEL_PATHS[dataset]
+    os.makedirs(cache_dir, exist_ok=True)
+    local_path = os.path.join(cache_dir, url.split("/")[-1])
+    if os.path.exists(local_path.replace('?dl=0', '')):
+        return local_path.replace('?dl=0', '')
+    headers = {'user-agent': 'Wget/1.16 (linux-gnu)'}
+    r = requests.get(url, stream=True, headers=headers)
+    size = int(r.headers.get("content-length", "0"))
+    with open(local_path, 'wb') as f:
+        pbar = tqdm(total=size, unit="iB", unit_scale=True)
+        for chunk in r.iter_content(chunk_size=chunk_size):
+            if chunk:
+                pbar.update(len(chunk))
+                f.write(chunk)
+    os.rename(local_path, local_path.replace('?dl=0', ''))
+    return local_path.replace('?dl=0', '')

requirements.txt

@@ -0,0 +1,2 @@
+git+https://github.com/jsu27/decomp_diffusion.git
+scikit-image==0.19.2