Update app.py

app.py

@@ -41,15 +41,19 @@ def pil_to_binary_mask(pil_image, threshold=0):
 # Duration timeを設定
 duration = 60
 
+device = "cuda"
+
 base_path = 'yisol/IDM-VTON'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
 
+# モデルのロードと初期化を関数外で行う
 unet = UNet2DConditionModel.from_pretrained(
     base_path,
     subfolder="unet",
     torch_dtype=torch.float16,
-)
+).to(device)
 unet.requires_grad_(False)
+
 tokenizer_one = AutoTokenizer.from_pretrained(
     base_path,
     subfolder="tokenizer",
@@ -68,38 +72,42 @@ text_encoder_one = CLIPTextModel.from_pretrained(
     base_path,
     subfolder="text_encoder",
     torch_dtype=torch.float16,
-)
+).to(device)
 text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
     base_path,
     subfolder="text_encoder_2",
     torch_dtype=torch.float16,
-)
+).to(device)
 image_encoder = CLIPVisionModelWithProjection.from_pretrained(
     base_path,
     subfolder="image_encoder",
     torch_dtype=torch.float16,
-)
+).to(device)
 vae = AutoencoderKL.from_pretrained(
     base_path,
     subfolder="vae",
     torch_dtype=torch.float16,
-)
+).to(device)
 
 UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
     base_path,
     subfolder="unet_encoder",
     torch_dtype=torch.float16,
-)
+).to(device)
 
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
+# モデルをGPUに転送
+parsing_model.model.to(device)
+openpose_model.preprocessor.body_estimation.model.to(device)
 UNet_Encoder.requires_grad_(False)
 image_encoder.requires_grad_(False)
 vae.requires_grad_(False)
 unet.requires_grad_(False)
 text_encoder_one.requires_grad_(False)
 text_encoder_two.requires_grad_(False)
+
 tensor_transfrom = transforms.Compose(
     [
         transforms.ToTensor(),
@@ -119,20 +127,16 @@ pipe = TryonPipeline.from_pretrained(
     scheduler=noise_scheduler,
     image_encoder=image_encoder,
     torch_dtype=torch.float16,
-)
+).to(device)
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU(duration=duration)  # duration変数を使用
 def start_tryon(
     dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, num_images
 ):
-    device = "cuda"
-
     start_time = time.time()  # 処理開始時間を記録
 
-    openpose_model.preprocessor.body_estimation.model.to(device)
-    pipe.to(device)
-    pipe.unet_encoder.to(device)
+    # device変数の再定義やモデルの.to(device)呼び出しを削除
 
     garm_img = garm_img.convert("RGB").resize((768, 1024))
     human_img_orig = dict["background"].convert("RGB")
@@ -355,4 +359,4 @@ with image_blocks as demo:
         api_name='tryon',
     )
 
-image_blocks.launch(show_error=True)
+image_blocks.launch(show_error=True)