Update app.py

app.py

@@ -5,6 +5,7 @@ import torchvision.transforms as transforms
 from torch.utils.data import Dataset, DataLoader
 import gradio as gr
 import sys
+import uuid
 import tqdm
 sys.path.append(os.path.abspath(os.path.join("", "..")))
 import gc
@@ -31,162 +32,102 @@ from diffusers import (
 from huggingface_hub import snapshot_download
 import spaces
 
+
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
 
-@spaces.GPU
-def load_models(device):
-    pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51" 
-    
-    revision = None
-    weight_dtype = torch.bfloat16
-    
-    # Load scheduler, tokenizer and models.
-    pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51", 
+device = "cuda"
+pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51" 
+revision = None
+weight_dtype = torch.bfloat16
+# Load scheduler, tokenizer and models.
+pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51", 
                                                 torch_dtype=torch.float16,safety_checker = None,
                                                 requires_safety_checker = False).to(device)
-    noise_scheduler = pipe.scheduler
-    del pipe
-    tokenizer = AutoTokenizer.from_pretrained(
+noise_scheduler = pipe.scheduler
+del pipe
+tokenizer = AutoTokenizer.from_pretrained(
             pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
         )
-    text_encoder = CLIPTextModel.from_pretrained(
+text_encoder = CLIPTextModel.from_pretrained(
             pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
         )
-    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
-    unet = UNet2DConditionModel.from_pretrained(
+vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
+unet = UNet2DConditionModel.from_pretrained(
             pretrained_model_name_or_path, subfolder="unet", revision=revision
         )
-    unet.requires_grad_(False)
-    unet.to(device, dtype=weight_dtype)
-    vae.requires_grad_(False)
-    
-    text_encoder.requires_grad_(False)
-    vae.requires_grad_(False)
-    vae.to(device, dtype=weight_dtype)
-    text_encoder.to(device, dtype=weight_dtype)
-    print("")
+unet.requires_grad_(False)
+unet.to(device, dtype=weight_dtype)
+vae.requires_grad_(False)
+    
+text_encoder.requires_grad_(False)
+vae.requires_grad_(False)
+vae.to(device, dtype=weight_dtype)
+text_encoder.to(device, dtype=weight_dtype)
+print("")
 
-    return unet, vae, text_encoder, tokenizer, noise_scheduler
 
-class main():
-    def __init__(self):
-        super(main, self).__init__()
-        
-        device = "cuda"
-        mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
-        std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
-        v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
-        proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
-        df = torch.load(f"{models_path}/files/identity_df.pt")
-        weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
-        pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+df = torch.load(f"{models_path}/files/identity_df.pt")
+weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
+pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 
-        self.device = device
-        self.mean = mean
-        self.std = std
-        self.v = v
-        self.proj = proj
-        self.df = df
-        self.weight_dimensions = weight_dimensions
-        self.pinverse = pinverse
-        
-        pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51" 
-    
-        revision = None
-        rank = 1
-        weight_dtype = torch.bfloat16
-        
-        # Load scheduler, tokenizer and models.
-        pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51", 
-                                                    torch_dtype=torch.float16,safety_checker = None,
-                                                    requires_safety_checker = False).to(device)
-        self.noise_scheduler = pipe.scheduler
-        del pipe
-        self.tokenizer = AutoTokenizer.from_pretrained(
-                pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
-            )
-        self.text_encoder = CLIPTextModel.from_pretrained(
-                pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
-            )
-        self.vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
-        self.unet = UNet2DConditionModel.from_pretrained(
-                pretrained_model_name_or_path, subfolder="unet", revision=revision
-            )
-        
-        self.unet.requires_grad_(False)
-        self.unet.to(device, dtype=weight_dtype)
-        self.vae.requires_grad_(False)
         
-        self.text_encoder.requires_grad_(False)
-        self.vae.requires_grad_(False)
-        self.vae.to(device, dtype=weight_dtype)
-        self.text_encoder.to(device, dtype=weight_dtype)
-        print("")
-
+young = get_direction(df, "Young", pinverse, 1000, device)
+young = debias(young, "Male", df, pinverse, device)
+young = debias(young, "Pointy_Nose", df, pinverse, device)
+young = debias(young, "Wavy_Hair", df, pinverse, device)
+young = debias(young, "Chubby", df, pinverse, device)
+young = debias(young, "No_Beard", df, pinverse, device)
+young = debias(young, "Mustache", df, pinverse, device)
+    
+pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
+pointy = debias(pointy, "Young", df, pinverse, device)
+pointy = debias(pointy, "Male", df, pinverse, device)
+pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
+pointy = debias(pointy, "Chubby", df, pinverse, device)
+pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
+        
+wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
+wavy = debias(wavy, "Young", df, pinverse, device)
+wavy = debias(wavy, "Male", df, pinverse, device)
+wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
+wavy = debias(wavy, "Chubby", df, pinverse, device)
+wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
+        
+thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
+thick = debias(thick, "Male", df, pinverse, device)
+thick = debias(thick, "Young", df, pinverse, device)
+thick = debias(thick, "Pointy_Nose", df, pinverse, device)
+thick = debias(thick, "Wavy_Hair", df, pinverse, device)
+thick = debias(thick, "Mustache", df, pinverse, device)
+thick = debias(thick, "No_Beard", df, pinverse, device)
+thick = debias(thick, "Sideburns", df, pinverse, device)
+thick = debias(thick, "Big_Nose", df, pinverse, device)
+thick = debias(thick, "Big_Lips", df, pinverse, device)
+thick = debias(thick, "Black_Hair", df, pinverse, device)
+thick = debias(thick, "Brown_Hair", df, pinverse, device)
+thick = debias(thick, "Pale_Skin", df, pinverse, device)
+thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
         
-        self.weights = None
-    
-        young = get_direction(df, "Young", pinverse, 1000, device)
-        young = debias(young, "Male", df, pinverse, device)
-        young = debias(young, "Pointy_Nose", df, pinverse, device)
-        young = debias(young, "Wavy_Hair", df, pinverse, device)
-        young = debias(young, "Chubby", df, pinverse, device)
-        young = debias(young, "No_Beard", df, pinverse, device)
-        young = debias(young, "Mustache", df, pinverse, device)
-        self.young = young
-    
-        pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
-        pointy = debias(pointy, "Young", df, pinverse, device)
-        pointy = debias(pointy, "Male", df, pinverse, device)
-        pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
-        pointy = debias(pointy, "Chubby", df, pinverse, device)
-        pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
-        self.pointy = pointy
-    
-        wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
-        wavy = debias(wavy, "Young", df, pinverse, device)
-        wavy = debias(wavy, "Male", df, pinverse, device)
-        wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
-        wavy = debias(wavy, "Chubby", df, pinverse, device)
-        wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
-        self.wavy = wavy
-    
-    
-        thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
-        thick = debias(thick, "Male", df, pinverse, device)
-        thick = debias(thick, "Young", df, pinverse, device)
-        thick = debias(thick, "Pointy_Nose", df, pinverse, device)
-        thick = debias(thick, "Wavy_Hair", df, pinverse, device)
-        thick = debias(thick, "Mustache", df, pinverse, device)
-        thick = debias(thick, "No_Beard", df, pinverse, device)
-        thick = debias(thick, "Sideburns", df, pinverse, device)
-        thick = debias(thick, "Big_Nose", df, pinverse, device)
-        thick = debias(thick, "Big_Lips", df, pinverse, device)
-        thick = debias(thick, "Black_Hair", df, pinverse, device)
-        thick = debias(thick, "Brown_Hair", df, pinverse, device)
-        thick = debias(thick, "Pale_Skin", df, pinverse, device)
-        thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
-        self.thick = thick
-
 
-   
-            
- 
-    @torch.no_grad()
-    @spaces.GPU(duration=120)
-    def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
-        device = self.device
-        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)
+@torch.no_grad()
+@spaces.GPU(duration=120)
+def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+    device = self.device
+    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(), 
+    network = LoRAw2w( self.weights.bfloat16(), self.mean.bfloat16(), self.std.bfloat16(), self.v[:, :1000].bfloat16(), 
                     self.unet,
                     rank=1,
                     multiplier=1.0,
@@ -196,68 +137,67 @@ class main():
         
 
         
-        generator = torch.Generator(device=device).manual_seed(seed)
-        latents = torch.randn(
+    generator = torch.Generator(device=device).manual_seed(seed)
+    latents = torch.randn(
                 (1, self.unet.in_channels, 512 // 8, 512 // 8),
                 generator = generator,
                 device = self.device
             ).bfloat16()
            
         
-        text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
         
-        text_embeddings = self.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 = self.tokenizer(
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = self.tokenizer(
                                     [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
                                 )
-        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
+    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)
+    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)
     
-            with network:
-                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
+        with network:
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
             
-            #guidance
-            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-            latents = self.noise_scheduler.step(noise_pred, t, latents).prev_sample
+        #guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        latents = self.noise_scheduler.step(noise_pred, t, latents).prev_sample
             
-        latents = 1 / 0.18215 * latents
-        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]
+    latents = 1 / 0.18215 * latents
+    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"))
+    image = Image.fromarray((image * 255).round().astype("uint8"))
         
-        return image
+    return image
 
 
-    @torch.no_grad()
-    @spaces.GPU(duration=120)
-    def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4): 
-        print("start")
-        device = self.device
-        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 = torch.load("model.pt").to(device)
-        self.young.to(device)
-        self.pointy.to(device)
-        self.wavy.to(device)
-        self.thick.to(device)
+@torch.no_grad()
+@spaces.GPU(duration=120)
+def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4): 
+    device = self.device
+    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 = torch.load("model.pt").to(device)
+    self.young.to(device)
+    self.pointy.to(device)
+    self.wavy.to(device)
+    self.thick.to(device)
 
-        network = LoRAw2w( self.weights.bfloat16(), self.mean.bfloat16(), self.std.bfloat16(), self.v[:, :1000].bfloat16(), 
+    network = LoRAw2w( self.weights.bfloat16(), self.mean.bfloat16(), self.std.bfloat16(), self.v[:, :1000].bfloat16(), 
                     self.unet,
                     rank=1,
                     multiplier=1.0,
@@ -266,90 +206,87 @@ class main():
                 ).to(device, torch.bfloat16)
         
         
-        original_weights = self.weights.clone()
+    original_weights = self.weights.clone()
             
-        #pad to same number of PCs
-        pcs_original = original_weights.shape[1]
-        pcs_edits = self.young.shape[1]
-        padding =  torch.zeros((1,pcs_original-pcs_edits)).to(device)
-        young_pad = torch.cat((self.young, padding), 1)
-        pointy_pad = torch.cat((self.pointy, padding), 1)
-        wavy_pad = torch.cat((self.wavy, padding), 1)
-        thick_pad = torch.cat((self.thick, padding), 1)
+    #pad to same number of PCs
+    pcs_original = original_weights.shape[1]
+    pcs_edits = self.young.shape[1]
+    padding =  torch.zeros((1,pcs_original-pcs_edits)).to(device)
+    young_pad = torch.cat((self.young, padding), 1)
+    pointy_pad = torch.cat((self.pointy, padding), 1)
+    wavy_pad = torch.cat((self.wavy, padding), 1)
+    thick_pad = torch.cat((self.thick, padding), 1)
             
         
-        edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*2e6*thick_pad
+    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(
+    generator = torch.Generator(device=device).manual_seed(seed)
+    latents = torch.randn(
                 (1, self.unet.in_channels, 512 // 8, 512 // 8),
                 generator = generator,
                 device = self.device
             ).bfloat16()
            
         
-        text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
         
-        text_embeddings = self.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 = self.tokenizer(
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = self.tokenizer(
                                     [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
                                 )
-        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
+    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)
+    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)
                 
-            if t>start_noise:
-                pass
-            elif t<=start_noise:
-                network.proj = torch.nn.Parameter(edited_weights)
-                network.reset()
+        if t>start_noise:
+            pass
+        elif t<=start_noise:
+            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
+        with network:
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
                     
                 
-            #guidance
-            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-            latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
+        #guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
             
-        latents = 1 / 0.18215 * latents
-        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"))
+    latents = 1 / 0.18215 * latents
+    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"))
 
-        
-        return image
+    return image
 
-    @torch.no_grad()
-    @spaces.GPU(duration=120)
-    def sample_then_run(self):
-        self.unet = UNet2DConditionModel.from_pretrained(
+@torch.no_grad()
+@spaces.GPU(duration=120)
+def sample_then_run(self):
+    self.unet = UNet2DConditionModel.from_pretrained(
             "stablediffusionapi/realistic-vision-v51" , subfolder="unet", revision=None
         )
-        self.unet.to(self.device, dtype=torch.bfloat16)
-        self.weights = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
-        
-        prompt = "sks person"
-        negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
-        seed = 5
-        cfg = 3.0
-        steps = 25
-        image = self.inference(prompt, negative_prompt, cfg, steps, seed)
-        torch.save(self.weights.cpu().detach(), "model.pt" )
-        return image, "model.pt"
+    self.unet.to(self.device, dtype=torch.bfloat16)
+    self.weights = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
+        
+    prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
+    seed = 5
+    cfg = 3.0
+    steps = 25
+    image = self.inference(prompt, negative_prompt, cfg, steps, seed)
+    torch.save(self.weights.cpu().detach(), "model.pt" )
+    return image, "model.pt"