fixed typo

app.py

@@ -118,7 +118,7 @@ def process_image(image, repaint, busyness):
     print("Saving masked image file to ", masked_img_file)
     image.save(masked_img_file)
     num = 64 # number of images to generate; we'll take the one with the most notes in the masked region
-    bs = numx
+    bs = num
     repaint = repaint
     seed_scale = 1.0
     CT_HOME = '.'

sample.py

@@ -5,9 +5,7 @@
 
 """Samples from k-diffusion models."""
 
-import gradio 
-import spaces
-import natten
+
 import argparse
 from pathlib import Path
 
@@ -24,11 +22,11 @@ from pom.v_diffusion import DDPM, LogSchedule, CrashSchedule
 #CHORD_BORDER = 8   # chord border size in pixels
 from pom.chords import CHORD_BORDER, img_batch_to_seq_emb, ChordSeqEncoder
 
+import spaces
 
 # ---- my mangled sampler that includes repaint 
 import torchsde 
 
-#@spaces.GPU
 class BatchedBrownianTree:
     """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
 
@@ -56,7 +54,6 @@ class BatchedBrownianTree:
         return w if self.batched else w[0]
 
 
-#@spaces.GPU
 class BrownianTreeNoiseSampler:
     """A noise sampler backed by a torchsde.BrownianTree.
 
@@ -94,7 +91,6 @@ def to_d(x, sigma, denoised):
     return (x - denoised) / append_dims(sigma, x.ndim)
 
 
-#@spaces.GPU
 @torch.no_grad()
 def my_sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., repaint=1):
     """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
@@ -129,7 +125,6 @@ def get_scalings(sigma, sigma_data=0.5):
     return c_skip, c_out, c_in
 
 
-#@spaces.GPU
 @torch.no_grad()
 def my_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, 
                     disable=None, eta=1., s_noise=1., noise_sampler=None, 
@@ -289,14 +284,12 @@ def sample(model, x, steps, eta, **extra_args):
 
 # Soft mask inpainting is just shrinking hard (binary) mask inpainting
 # Given a float-valued soft mask (values between 0 and 1), get the binary mask for this particular step
-#@spaces.GPU
 def get_bmask(i, steps, mask):
     strength = (i+1)/(steps)
     # convert to binary mask
     bmask = torch.where(mask<=strength,1,0)
     return bmask
 
-#@spaces.GPU
 def make_cond_model_fn(model, cond_fn):
     def cond_model_fn(x, sigma, **kwargs):
         with torch.enable_grad():
@@ -312,7 +305,6 @@ def make_cond_model_fn(model, cond_fn):
 # For sampling, set both init_data and mask to None
 # For variations, set init_data 
 # For inpainting, set both init_data & mask 
-#@spaces.GPU
 def sample_k(
         model_fn, 
         noise, 
@@ -425,7 +417,7 @@ def infer_mask_from_init_img(img, mask_with='white'):
         mask[img[2,:,:]==1] = 1  # blue
     return mask*1.0
 
-#@spaces.GPU
+
 def grow_mask(init_mask, grow_by=2):
     "adds a border of grow_by pixels to the mask, by growing it grow_by times. If grow_by=0, does nothing"
     new_mask = init_mask.clone()
@@ -434,7 +426,7 @@ def grow_mask(init_mask, grow_by=2):
         new_mask[1:-1,1:-1] = (new_mask[1:-1,1:-1] + new_mask[0:-2,1:-1] + new_mask[2:,1:-1] + new_mask[1:-1,0:-2] + new_mask[1:-1,2:]) > 0 
     return new_mask
 
-#@spaces.GPU
+
 def add_seeding(init_image, init_mask, grow_by=0, seed_scale=1.0):
     "adds extra noise inside mask"
     init_mask = grow_mask(init_mask, grow_by=grow_by)  # make the mask bigger
@@ -448,15 +440,13 @@ def add_seeding(init_image, init_mask, grow_by=0, seed_scale=1.0):
     init_image[2,:,:] = init_image[2,:,:] * (1-init_mask) - 1.0*init_mask
     return init_image
 
-#@spaces.GPU
+
 def get_init_image_and_mask(args, device):
     convert_tensor = transforms.ToTensor()
     init_image = Image.open(args.init_image).convert('RGB')
     init_image = convert_tensor(init_image)
     #normalize image from 0..1 to -1..1
     init_image = (2.0 * init_image) - 1.0
-
-
     init_mask = torch.ones(init_image.shape[-2:])  # ones are where stuff will change, zeros will stay the same
 
     inpaint_task = 'infer'  # infer mask from init_image
@@ -522,7 +512,115 @@ def get_init_image_and_mask(args, device):
     init_mask = init_mask.unsqueeze(0).unsqueeze(1).repeat(args.batch_size,3,1,1).float()
     return init_image.to(device), init_mask.to(device)
 
-#@spaces.GPU # generates an error
+
+
+
+# wrapper compatible with ZeroGPU, callable from outside
+@spaces.GPU
+def zero_wrapper(args, device): 
+    global init_image, init_mask
+
+    config = K.config.load_config(args.config if args.config else args.checkpoint)
+    model_config = config['model']
+    # TODO: allow non-square input sizes
+    assert len(model_config['input_size']) == 2 and model_config['input_size'][0] == model_config['input_size'][1]
+    size = model_config['input_size']
+
+    print('zero_wrapper: Using device:', device, flush=True)
+
+    inner_model = K.config.make_model(config).eval().requires_grad_(False).to(device)
+    cse = None # ChordSeqEncoder().eval().requires_grad_(False).to(device)  # add chord embedding-maker to main model
+    if cse is not None:
+        inner_model.cse = cse
+    try:
+        inner_model.load_state_dict(safetorch.load_file(args.checkpoint))
+    except:
+        #ckpt = torch.load(args.checkpoint).to(device)
+        ckpt = torch.load(args.checkpoint, map_location='cpu')
+        inner_model.load_state_dict(ckpt['model'])
+
+    print('Parameters:', K.utils.n_params(inner_model))
+    model = K.Denoiser(inner_model, sigma_data=model_config['sigma_data'])
+
+    sigma_min = model_config['sigma_min']
+    sigma_max = model_config['sigma_max']
+    torch.set_float32_matmul_precision('high')
+    extra_args = {}
+    init_image, init_mask = None, None
+    if args.init_image is not None:
+        init_image, init_mask = get_init_image_and_mask(args, device)
+        init_image = init_image.to(device)
+        init_mask = init_mask.to(device)
+    @torch.no_grad()
+    @K.utils.eval_mode(model)
+    def run():
+        global init_image, init_mask
+        if accelerator.is_local_main_process:
+            tqdm.write('Sampling...')
+        sigmas = K.sampling.get_sigmas_karras(args.steps, sigma_min, sigma_max, rho=7., device=device)
+
+        def sample_fn(n, debug=True):
+            x = torch.randn([n, model_config['input_channels'], size[0], size[1]], device=device) * sigma_max
+            print("n, sigma_max, x.min, x.max = ", n, sigma_max, x.min(), x.max())
+
+            if args.init_image is not None:
+                init_data, mask = get_init_image_and_mask(args, device)
+                init_data = args.seed_scale*x*mask + (1-mask)*init_data  # extra nucleation?
+                if cse is not None: 
+                    chord_cond = img_batch_to_seq_emb(init_data, inner_model.cse).to(device)
+                else: 
+                    chord_cond = None
+                #print("init_data.shape, init_data.min, init_data.max = ", init_data.shape, init_data.min(), init_data.max())
+            else:
+                init_data, mask, chord_cond = None, None, None
+            # chord_cond doesn't work anyway so f it: 
+            chord_cond = None
+
+            print("chord_cond = ", chord_cond)
+            if chord_cond is not None: 
+                extra_args['chord_cond'] = chord_cond
+            # these two work:
+            #x_0 = K.sampling.sample_lms(model, x, sigmas, disable=not accelerator.is_local_main_process, extra_args=extra_args)
+            #x_0 = K.sampling.sample_dpmpp_2m_sde(model, x, sigmas, disable=not accelerator.is_local_main_process, extra_args=extra_args)
+
+            noise = torch.randn([n, model_config['input_channels'], size[0], size[1]], device=device) 
+
+            sampler_type="my-dpmpp-2m-sde"  # "k-lms"
+            #sampler_type="my-sample-euler"
+            #sampler_type="dpmpp-2m-sde"  
+            #sampler_type = "dpmpp-3m-sde"
+            #sampler_type = "k-dpmpp-2s-ancestral"
+            print("dtypes:", [x.dtype if x is not None else None  for x in [noise, init_data, mask, chord_cond]])
+            x_0 = sample_k(inner_model, noise, sampler_type=sampler_type, 
+                           init_data=init_data, mask=mask, steps=args.steps, 
+                           sigma_min=sigma_min, sigma_max=sigma_max, rho=7., 
+                           device=device, model_config=model_config, repaint=args.repaint, 
+                           **extra_args)
+            #x_0 = sample_k(inner_model, noise, sampler_type="dpmpp-2m-sde", steps=100,  sigma_min=0.5, sigma_max=50, rho=1., device=device,  model_config=model_config, **extra_args)
+            print("x_0.min, x_0.max = ", x_0.min(), x_0.max())
+            if x_0.isnan().any():
+                assert False, "x_0 has NaNs"
+            
+            # do gpu garbage collection before proceeding
+            torch.cuda.empty_cache()
+            return x_0
+        
+        x_0 = K.evaluation.compute_features(accelerator, sample_fn, lambda x: x, args.n, args.batch_size)
+        if accelerator.is_main_process:
+            for i, out in enumerate(x_0):
+                filename = f'{args.prefix}_{i:05}.png'
+                K.utils.to_pil_image(out).save(filename)
+
+    try:
+        run()
+    except KeyboardInterrupt:
+        pass
+
+
+
+
+
+
 def main():
     global init_image, init_mask
     p = argparse.ArgumentParser(description=__doc__,
@@ -574,12 +672,7 @@ def main():
     sigma_min = model_config['sigma_min']
     sigma_max = model_config['sigma_max']
 
-    # SHH modified
     torch.set_float32_matmul_precision('high')
-    #class_cond = torch.tensor([0]).to(device)
-    #num_classes = 10
-    #class_cond = torch.remainder(torch.arange(0, args.n), num_classes).int().to(device)
-    #extra_args = {'class_cond':class_cond}
     extra_args = {}
     init_image, init_mask = None, None
     if args.init_image is not None:
@@ -595,11 +688,6 @@ def main():
             tqdm.write('Sampling...')
         sigmas = K.sampling.get_sigmas_karras(args.steps, sigma_min, sigma_max, rho=7., device=device)
 
-        #ddpm_sampler = DDPM(model)
-        #model_fn = model
-        #ddpm_sampler = K.external.VDenoiser(model_fn)
-
-        #@spaces.GPU
         def sample_fn(n, debug=True):
             x = torch.randn([n, model_config['input_channels'], size[0], size[1]], device=device) * sigma_max
             print("n, sigma_max, x.min, x.max = ", n, sigma_max, x.min(), x.max())