adding zerogpu decorators

sample.py

@@ -1,9 +1,13 @@
 #!/usr/bin/env python3
 
 # Code by Kat Crowson in k-diffusion repo, modified by Scott H Hawley (SHH)
+# Modified by Scott H. Hawley for masking, ZeroGPU ets.
 
 """Samples from k-diffusion models."""
 
+import gradio 
+import spaces
+import natten
 import argparse
 from pathlib import Path
 
@@ -24,6 +28,7 @@ from pom.chords import CHORD_BORDER, img_batch_to_seq_emb, ChordSeqEncoder
 # ---- my mangled sampler that includes repaint 
 import torchsde 
 
+@spaces.GPU
 class BatchedBrownianTree:
     """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
 
@@ -51,6 +56,7 @@ class BatchedBrownianTree:
         return w if self.batched else w[0]
 
 
+@spaces.GPU
 class BrownianTreeNoiseSampler:
     """A noise sampler backed by a torchsde.BrownianTree.
 
@@ -88,6 +94,7 @@ 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)."""
@@ -122,6 +129,7 @@ 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, 
@@ -281,12 +289,14 @@ 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():
@@ -302,6 +312,7 @@ 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, 
@@ -399,6 +410,7 @@ def sample_k(
 
 
 ## ---- end stable-audio-tools
+@spaces.GPU
 def infer_mask_from_init_img(img, mask_with='white'):
     """given an image with mask areas marked, extract the mask itself
        note, this works whether image is normalized on 0..1 or -1..1, but not 0..255"""
@@ -413,6 +425,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()
@@ -421,7 +434,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
@@ -435,7 +448,7 @@ 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')
@@ -509,7 +522,7 @@ 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
 def main():
     global init_image, init_mask
     p = argparse.ArgumentParser(description=__doc__,
@@ -586,6 +599,7 @@ def main():
         #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())