fix zero

app.py

@@ -5,35 +5,29 @@ import torch
 import numpy as np
 import torch
 from pytorch_lightning import seed_everything
-from model_lib.modules import MoMA_main_modal
 from model_lib.utils import parse_args
 import os
 os.environ["CUDA_VISIBLE_DEVICES"]="0"
 
 title = "MoMA"
-description = "This model has to run on GPU. By default, we load the model with 4-bit quantization to make it fit in smaller hardware."
+description = "This model has to run on GPU. Please find our project page at https://moma-adapter.github.io/."
 device = torch.device('cuda')
-
 seed_everything(0)
 args = parse_args()
-#load MoMA from HuggingFace. Auto download
-model = MoMA_main_modal(args).to(device, dtype=torch.float16)
 
-@spaces.GPU
 def MoMA_demo(rgb, subject, prompt, strength, seed):
-    with torch.no_grad():
-        generated_image = model.generate_images(rgb, subject, prompt, strength=strength, seed=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
 
 @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
 
-
 gr.Interface(
     inference,
     [gr.Image(type="pil", label="Input RGB"),

model_lib/moMA_generator.py

@@ -155,7 +155,6 @@ class MoMA_generator:
         return image_prompt_embeds, uncond_image_prompt_embeds
 
     # feature are from self-attention layers of Unet: feed reference image to Unet with t=0
-    @spaces.GPU
     def get_image_selfAttn_feature(
             self,
             pil_image,

model_lib/modules.py

@@ -112,7 +112,6 @@ class MoMA_main_modal(nn.Module):
             module.train = False
             module.requires_grad_(False)
 
-    @spaces.GPU
     def forward_MLLM(self,batch):
         llava_processeds,subjects,prompts = batch['llava_processed'].half().to(self.device),batch['label'],batch['text']
         
@@ -138,6 +137,7 @@ 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)