Update app.py

app.py

@@ -243,7 +243,25 @@ class main():
     @spaces.GPU
     def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4): 
         device = self.device
-        original_weights = self,network.proj.clone()
+        self.unet.to(device)
+        self.text_encoder.to(device)
+        self.vae.to(device)
+        self.mean.to(device)
+        self.std.to(device)
+        self.v.to(device)
+        self.proj.to(device)
+        self.weights.to(device)
+
+        network = LoRAw2w( self.weights.bfloat16(), self.mean.bfloat16(), self.std.bfloat16(), self.v[:, :1000].bfloat16(), 
+                    self.unet,
+                    rank=1,
+                    multiplier=1.0,
+                    alpha=27.0,
+                    train_method="xattn-strict"
+                ).to(device, torch.bfloat16)
+        
+        
+        original_weights = self.weights.clone()
             
         #pad to same number of PCs
         pcs_original = original_weights.shape[1]
@@ -256,7 +274,7 @@ class main():
             
         
         edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*2e6*thick_pad
-           
+
         generator = torch.Generator(device=device).manual_seed(seed)
         latents = torch.randn(
                 (1, self.unet.in_channels, 512 // 8, 512 // 8),
@@ -267,19 +285,19 @@ class main():
         
         text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
         
-        text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
+        text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
         
         max_length = text_input.input_ids.shape[-1]
-        uncond_input = tokenizer(
+        uncond_input = self.tokenizer(
                                     [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
                                 )
-        uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
-        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
-        noise_scheduler.set_timesteps(ddim_steps) 
-        latents = latents * noise_scheduler.init_noise_sigma
-            
+        uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(device))[0]
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).bfloat16()
+        self.noise_scheduler.set_timesteps(ddim_steps) 
+        latents = latents * self.noise_scheduler.init_noise_sigma
+           
         
-         
+            
         for i,t in enumerate(tqdm.tqdm(self.noise_scheduler.timesteps)):
             latent_model_input = torch.cat([latents] * 2)
             latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep=t)
@@ -287,11 +305,10 @@ class main():
             if t>start_noise:
                 pass
             elif t<=start_noise:
-                self.network.proj = torch.nn.Parameter(edited_weights)
-                self.network.reset()
-        
-        
-            with self.network:
+                network.proj = torch.nn.Parameter(edited_weights)
+                network.reset()
+          
+            with network:
                 noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
                     
                 
@@ -301,16 +318,13 @@ class main():
             latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
             
         latents = 1 / 0.18215 * latents
-        image = self.vae.decode(latents).sample
+        image = self.vae.decode(latents.float()).sample
         image = (image / 2 + 0.5).clamp(0, 1)
         
         image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
         
         image = Image.fromarray((image * 255).round().astype("uint8"))
-        
-        #reset weights back to original 
-        self.network.proj = torch.nn.Parameter(original_weights)
-        self.network.reset()
+
         
         return image