float16

app.py

@@ -28,7 +28,7 @@ def inference(rgb, mask, subject, prompt):
 seed_everything(0)
 args = parse_args()
 #load MoMA from HuggingFace. Auto download
-model = MoMA_main_modal(args).to(args.device, dtype=torch.bfloat16)
+model = MoMA_main_modal(args).to(args.device, dtype=torch.float16)
 
 
 ################ change texture ##################

model_lib/moMA_generator.py

@@ -93,7 +93,7 @@ class MoMA_generator:
         print('Loading StableDiffusion: Realistic_Vision...')
         self.pipe = StableDiffusionPipeline.from_pretrained(
             "SG161222/Realistic_Vision_V4.0_noVAE",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float16,
             scheduler=noise_scheduler,
             vae=vae,
             feature_extractor=None,
@@ -112,7 +112,7 @@ class MoMA_generator:
             cross_attention_dim=768,
             clip_embeddings_dim=1024,
             clip_extra_context_tokens=4,
-        ).to(self.device, dtype=torch.bfloat16)
+        ).to(self.device, dtype=torch.float16)
         return image_proj_model
         
     def set_ip_adapter(self):

model_lib/modules.py

@@ -87,7 +87,7 @@ class MoMA_main_modal(nn.Module):
         
         add_function(self.model_llava)
 
-        self.mapping = LlamaMLP_mapping(4096,1024).to(self.device, dtype=torch.bfloat16)
+        self.mapping = LlamaMLP_mapping(4096,1024).to(self.device, dtype=torch.float16)
         self.load_saved_components()
         self.freeze_modules()
 
@@ -140,7 +140,7 @@ class MoMA_main_modal(nn.Module):
         batch = Dataset_evaluate_MoMA(rgb_path, prompt, subject, mask_path,self)
         self.moMA_generator.set_selfAttn_strength(strength)
         
-        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True):
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True):
             with torch.no_grad(): 
                 ### key steps
                 llava_emb = self.forward_MLLM(batch).clone().detach()