Update app.py

app.py

@@ -290,29 +290,29 @@ def sample_then_run(self):
 
 
 
-    class CustomImageDataset(Dataset):
-        def __init__(self, images, transform=None):
-            self.images = images
-            self.transform = transform
-        
-        def __len__(self):
-            return len(self.images)
-        
-        def __getitem__(self, idx):
-            image = self.images[idx]
-            if self.transform:
-                image = self.transform(image)
-            return image
-        
-    @spaces.GPU
-    def invert(self, image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
+class CustomImageDataset(Dataset):
+    def __init__(self, images, transform=None):
+        self.images = images
+        self.transform = transform
+        
+    def __len__(self):
+        return len(self.images)
+    
+    def __getitem__(self, idx):
+        image = self.images[idx]
+        if self.transform:
+            image = self.transform(image)
+        return image
+        
+@spaces.GPU
+def invert(self, image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
             
-        del unet
-        del network
-        unet, _, _, _, _ = load_models(device)
+    del unet
+    del network
+    unet, _, _, _, _ = load_models(device)
             
-        proj = torch.zeros(1,pcs).bfloat16().to(device)
-        network = LoRAw2w( proj, mean, std, v[:, :pcs], 
+    proj = torch.zeros(1,pcs).bfloat16().to(device)
+    network = LoRAw2w( proj, mean, std, v[:, :pcs], 
                                 unet,
                                 rank=1,
                                 multiplier=1.0,
@@ -320,87 +320,83 @@ def sample_then_run(self):
                                 train_method="xattn-strict"
                             ).to(device, torch.bfloat16)
     
-        ### load mask
-        mask = transforms.Resize((64,64), interpolation=transforms.InterpolationMode.BILINEAR)(mask)
-        mask = torchvision.transforms.functional.pil_to_tensor(mask).unsqueeze(0).to(device).bfloat16()[:,0,:,:].unsqueeze(1)
-        ### check if an actual mask was draw, otherwise mask is just all ones
-        if torch.sum(mask) == 0:
-            mask = torch.ones((1,1,64,64)).to(device).bfloat16()
+    ### load mask
+    mask = transforms.Resize((64,64), interpolation=transforms.InterpolationMode.BILINEAR)(mask)
+    mask = torchvision.transforms.functional.pil_to_tensor(mask).unsqueeze(0).to(device).bfloat16()[:,0,:,:].unsqueeze(1)
+    ### check if an actual mask was draw, otherwise mask is just all ones
+    if torch.sum(mask) == 0:
+        mask = torch.ones((1,1,64,64)).to(device).bfloat16()
                 
-        ### single image dataset
-        image_transforms = transforms.Compose([transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
+    ### single image dataset
+    image_transforms = transforms.Compose([transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
                                                         transforms.RandomCrop(512),
                                                         transforms.ToTensor(),
                                                         transforms.Normalize([0.5], [0.5])])
         
         
-        train_dataset = CustomImageDataset(image, transform=image_transforms)
-        train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True) 
+    train_dataset = CustomImageDataset(image, transform=image_transforms)
+    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True) 
         
-        ### optimizer 
-        optim = torch.optim.Adam(network.parameters(), lr=lr, weight_decay=weight_decay)    
+    ### optimizer 
+    optim = torch.optim.Adam(network.parameters(), lr=lr, weight_decay=weight_decay)    
         
-        ### training loop
-        unet.train()
-        for epoch in tqdm.tqdm(range(epochs)):
-            for batch in train_dataloader:
-                ### prepare inputs
-                batch = batch.to(device).bfloat16()
-                latents = vae.encode(batch).latent_dist.sample()
-                latents = latents*0.18215
-                noise = torch.randn_like(latents)
-                bsz = latents.shape[0]
+    ### training loop
+    unet.train()
+    for epoch in tqdm.tqdm(range(epochs)):
+        for batch in train_dataloader:
+            ### prepare inputs
+            batch = batch.to(device).bfloat16()
+            latents = vae.encode(batch).latent_dist.sample()
+            latents = latents*0.18215
+            noise = torch.randn_like(latents)
+            bsz = latents.shape[0]
                  
-                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
-                timesteps = timesteps.long()
-                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
-                text_input = tokenizer("sks person", padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
-                text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
-        
-                ### loss + sgd step
-                with network:
-                    model_pred = unet(noisy_latents, timesteps, text_embeddings).sample
-                    loss = torch.nn.functional.mse_loss(mask*model_pred.float(), mask*noise.float(), reduction="mean")
-                    optim.zero_grad()
-                    loss.backward()
-                    optim.step()
-        
-        ### return optimized network
-        return network
+            timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+            timesteps = timesteps.long()
+            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+            text_input = tokenizer("sks person", padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+            text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
+        
+            ### loss + sgd step
+            with network:
+                model_pred = unet(noisy_latents, timesteps, text_embeddings).sample
+                loss = torch.nn.functional.mse_loss(mask*model_pred.float(), mask*noise.float(), reduction="mean")
+                optim.zero_grad()
+                loss.backward()
+                optim.step()
+        
+    ### return optimized network
+    return network
 
 
-    @spaces.GPU
-    def run_inversion(self, dict, pcs, epochs, weight_decay,lr):
-        init_image = dict["image"].convert("RGB").resize((512, 512))
-        mask = dict["mask"].convert("RGB").resize((512, 512))
-        network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
-        
-        
-        #sample an image
-        prompt = "sks person"
-        negative_prompt = "low quality, blurry, unfinished, nudity"
-        seed = 5
-        cfg = 3.0
-        steps = 25
-        image = inference( prompt, negative_prompt, cfg, steps, seed)
-        torch.save(network.proj, "model.pt" )
-        return image, "model.pt"
+@spaces.GPU
+def run_inversion(self, dict, pcs, epochs, weight_decay,lr):
+    init_image = dict["image"].convert("RGB").resize((512, 512))
+    mask = dict["mask"].convert("RGB").resize((512, 512))
+    network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
+    #sample an image
+    prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
+    seed = 5
+    cfg = 3.0
+    steps = 25
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    torch.save(network.proj, "model.pt" )
+    return image, "model.pt"
     
     
-    @spaces.GPU
-    def file_upload(self, file):
-          
-        proj = torch.load(file.name).to(device)
-        
-        #pad to 10000 Principal components to keep everything consistent
-        pcs = proj.shape[1]
-        padding =  torch.zeros((1,10000-pcs)).to(device)
-        proj = torch.cat((proj, padding), 1)
+@spaces.GPU
+def file_upload(self, file):
+    proj = torch.load(file.name).to(device)
         
-        unet, _, _, _, _ = load_models(device)
+    #pad to 10000 Principal components to keep everything consistent
+    pcs = proj.shape[1]
+    padding =  torch.zeros((1,10000-pcs)).to(device)
+    proj = torch.cat((proj, padding), 1)
+    unet, _, _, _, _ = load_models(device)
         
         
-        network = LoRAw2w( proj, mean, std, v[:, :10000], 
+    network = LoRAw2w( proj, mean, std, v[:, :10000], 
                                 unet,
                                 rank=1,
                                 multiplier=1.0,
@@ -409,13 +405,13 @@ def sample_then_run(self):
                             ).to(device, torch.bfloat16)
                 
             
-        prompt = "sks person"
-        negative_prompt = "low quality, blurry, unfinished, nudity"
-        seed = 5
-        cfg = 3.0
-        steps = 25
-        image = inference( prompt, negative_prompt, cfg, steps, seed)
-        return image
+    prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
+    seed = 5
+    cfg = 3.0
+    steps = 25
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    return image