fix zero

app.py

@@ -15,18 +15,15 @@ device = torch.device('cuda')
 seed_everything(0)
 args = parse_args()
 
-def MoMA_demo(rgb, subject, prompt, strength, seed):
-    from model_lib.modules import MoMA_main_modal
-    model = MoMA_main_modal(args).to(device, dtype=torch.float16)
-    generated_image = model.generate_images(rgb, subject, prompt, strength=strength, seed=seed)
-    return generated_image
+from model_lib.modules import MoMA_main_modal
+model = MoMA_main_modal(args).to(device, dtype=torch.float16)
 
 @spaces.GPU
 def inference(rgb, subject, prompt, strength, seed):
     seed = int(seed) if seed else 0
     seed = seed if not seed == 0 else np.random.randint(0,1000)
-    result = MoMA_demo(rgb, subject, prompt, strength, seed)
-    return result
+    generated_image = model.generate_images(rgb, subject, prompt, strength=strength, seed=seed)
+    return generated_image
 
 gr.Interface(
     inference,

model_lib/moMA_generator.py

@@ -1,6 +1,3 @@
-import spaces
-
-
 import torch
 from model_lib.attention_processor import IPAttnProcessor, IPAttnProcessor_Self, get_mask_from_cross
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
@@ -98,7 +95,7 @@ class MoMA_generator:
             vae=vae,
             feature_extractor=None,
             safety_checker=None,
-        )
+        ).to(self.device)
 
         self.unet = self.pipe.unet
         add_function(self.pipe)
@@ -112,7 +109,7 @@ class MoMA_generator:
             cross_attention_dim=768,
             clip_embeddings_dim=1024,
             clip_extra_context_tokens=4,
-        )
+        ).to(self.device, dtype=torch.float16)
         return image_proj_model
         
     def set_ip_adapter(self):
@@ -129,9 +126,9 @@ class MoMA_generator:
                 block_id = int(name[len("down_blocks.")])
                 hidden_size = unet.config.block_out_channels[block_id]
             if cross_attention_dim is None:
-                attn_procs[name] = IPAttnProcessor_Self(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4)
+                attn_procs[name] = IPAttnProcessor_Self(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4).to(self.device, dtype=torch.float16)
             else:
-                attn_procs[name] = IPAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4)
+                attn_procs[name] = IPAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4).to(self.device, dtype=torch.float16)
         unet.set_attn_processor(attn_procs)
 
     @torch.inference_mode()

model_lib/modules.py

@@ -1,5 +1,3 @@
-import spaces
-
 import os
 import torch
 import torch.nn as nn
@@ -84,11 +82,11 @@ class MoMA_main_modal(nn.Module):
         
         print('Loading MoMA: its Multi-modal LLM...')
         model_name = get_model_name_from_path(args.model_path)
-        self.tokenizer_llava, self.model_llava, self.image_processor_llava, self.context_len_llava = load_pretrained_model(args.model_path, None, model_name, load_8bit=self.args.load_8bit, load_4bit=self.args.load_4bit)
+        self.tokenizer_llava, self.model_llava, self.image_processor_llava, self.context_len_llava = load_pretrained_model(args.model_path, None, model_name, load_8bit=self.args.load_8bit, load_4bit=self.args.load_4bit, device=args.device)
         
         add_function(self.model_llava)
 
-        self.mapping = LlamaMLP_mapping(4096,1024)
+        self.mapping = LlamaMLP_mapping(4096,1024).to(self.device, dtype=torch.float16)
         self.load_saved_components()
         self.freeze_modules()
 
@@ -137,7 +135,6 @@ class MoMA_main_modal(nn.Module):
     def reset(self):
         self.moMA_generator.reset_all()
 
-    @torch.no_grad()
     def generate_images(self, rgb_path, subject, prompt, strength=1.0, num=1, seed=0):
         batch = Dataset_evaluate_MoMA(rgb_path, prompt, subject,self)
         self.moMA_generator.set_selfAttn_strength(strength)

model_lib/utils.py

@@ -10,7 +10,7 @@ def parse_args():
     parser.add_argument("--model_path",type=str,default="KunpengSong/MoMA_llava_7b",help="fine tuned llava (Multi-modal LLM decoder)")
     args = parser.parse_known_args()[0]
     args.device = torch.device("cuda", 0)
-    args.load_8bit, args.load_4bit = False, False
+    args.load_8bit, args.load_4bit = False, True
     return args
 
 def show_PIL_image(tensor):