fix zero

app.py

@@ -19,6 +19,7 @@ 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)

model_lib/moMA_generator.py

@@ -1,3 +1,6 @@
+import spaces
+
+
 import torch
 from model_lib.attention_processor import IPAttnProcessor, IPAttnProcessor_Self, get_mask_from_cross
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
@@ -109,7 +112,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):
@@ -126,9 +129,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).to(self.device, dtype=torch.float16)
+                attn_procs[name] = IPAttnProcessor_Self(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4)
             else:
-                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)
+                attn_procs[name] = IPAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim,scale=1.0,num_tokens=4)
         unet.set_attn_processor(attn_procs)
 
     @torch.inference_mode()
@@ -152,6 +155,7 @@ 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

@@ -1,3 +1,5 @@
+import spaces
+
 import os
 import torch
 import torch.nn as nn
@@ -110,6 +112,7 @@ 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']