update

infer_api.py

@@ -171,10 +171,11 @@ def process_image(image, totensor, width, height):
 @spaces.GPU
 @torch.no_grad()
 def inference(validation_pipeline, bkg_remover, input_image, vae, feature_extractor, image_encoder, unet, ref_unet, tokenizer,
-              text_encoder, pretrained_model_path, generator, validation, val_width, val_height, unet_condition_type,
+              text_encoder, pretrained_model_path, validation, val_width, val_height, unet_condition_type,
               use_noise=True, noise_d=256, crop=False, seed=100, timestep=20):
     set_seed(seed)
-
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
     totensor = transforms.ToTensor()
 
     prompts = "high quality, best quality"
@@ -278,10 +279,6 @@ def run_multiview_infer(data, pipeline, cfg: TestConfig, num_levels=3):
         generator = None
     else:
         generator = torch.Generator(device=pipeline.unet.device).manual_seed(cfg.seed)
-
-    if torch.cuda.is_available():
-        pipeline.unet.enable_xformers_memory_efficient_attention()
-        pipeline.to(device)
     
     images_cond = []
     results = {}
@@ -341,11 +338,14 @@ def run_multiview_infer(data, pipeline, cfg: TestConfig, num_levels=3):
     torch.cuda.empty_cache()    
     return results
 
-
+@spaces.GPU
 def load_multiview_pipeline(cfg):
     pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
         cfg.pretrained_path,
         torch_dtype=torch.float16,)
+    pipeline.unet.enable_xformers_memory_efficient_attention()
+    if torch.cuda.is_available():
+        pipeline.to(device)
     return pipeline
 
 
@@ -870,10 +870,9 @@ class InferCanonicalAPI:
         self.bkg_remover = BkgRemover()
 
     def canonicalize(self, image, seed):
-        generator = torch.Generator(device=device).manual_seed(seed)
         return inference(
             self.validation_pipeline, self.bkg_remover, image, self.vae, self.feature_extractor, self.image_encoder, self.unet, self.ref_unet, self.tokenizer, self.text_encoder,
-            self.pretrained_model_path, generator, self.validation, self.width_input, self.height_input, self.unet_condition_type,
+            self.pretrained_model_path, self.validation, self.width_input, self.height_input, self.unet_condition_type,
             use_noise=self.use_noise, noise_d=self.noise_d, crop=True, seed=seed, timestep=self.timestep
         )