fix zero

app.py

@@ -1,12 +1,10 @@
 import spaces
+
 import gradio as gr
-import cv2
 import torch
 import numpy as np
-from torchvision import transforms
 import torch
 from pytorch_lightning import seed_everything
-from torchvision.utils import save_image
 from model_lib.modules import MoMA_main_modal
 from model_lib.utils import parse_args
 import os
@@ -14,11 +12,12 @@ 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."
+device = torch.device('cuda')
 
 seed_everything(0)
 args = parse_args()
 #load MoMA from HuggingFace. Auto download
-model = MoMA_main_modal(args).to(args.device, dtype=torch.float16)
+model = MoMA_main_modal(args).to(device, dtype=torch.float16)
 
 def MoMA_demo(rgb, subject, prompt, strength, seed):
     with torch.no_grad():

dataset_lib/dataset_eval_MoMA.py

@@ -2,7 +2,7 @@ from PIL import Image
 import numpy as np
 import torch
 from torchvision import transforms
-from llava.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+from llava.mm_utils import process_images
 from rembg import remove
 
 def create_binary_mask(image):

model_lib/attention_processor.py

@@ -4,8 +4,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
 import math
-from torchvision.utils import save_image
-import torchvision.transforms as T
 
 def get_mask_from_cross(attn_processors):
     reference_masks = []

model_lib/moMA_generator.py

@@ -1,7 +1,4 @@
-from typing import List
 import torch
-from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
-from PIL import Image
 from model_lib.attention_processor import IPAttnProcessor, IPAttnProcessor_Self, get_mask_from_cross
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
 import tqdm

model_lib/modules.py

@@ -1,5 +1,4 @@
 import os
-from PIL import Image
 import torch
 import torch.nn as nn
 from typing import List, Optional