Update app.py

app.py

@@ -23,7 +23,7 @@ from preprocess.humanparsing.run_parsing import Parsing
 from preprocess.openpose.run_openpose import OpenPose
 from detectron2.data.detection_utils import convert_PIL_to_numpy, _apply_exif_orientation
 from torchvision.transforms.functional import to_pil_image
-import time  # timeモジュールをインポート
+
 
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
@@ -38,22 +38,16 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
-# 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",
@@ -72,43 +66,38 @@ 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)
-parsing_model.parsenet.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(),
@@ -128,16 +117,19 @@ 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変数を使用
+
+@spaces.GPU(duration=60)  # 実行時間を60秒に設定
 def start_tryon(
     dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, num_images
 ):
-    start_time = time.time()  # 処理開始時間を記録
+    device = "cuda"
 
-    # device変数の再定義やモデルの.to(device)呼び出しを削除
+    openpose_model.preprocessor.body_estimation.model.to(device)
+    pipe.to(device)
+    pipe.unet_encoder.to(device)
 
     garm_img = garm_img.convert("RGB").resize((768, 1024))
     human_img_orig = dict["background"].convert("RGB")
@@ -226,11 +218,6 @@ def start_tryon(
     yield output_images.value, mask_gray
 
     for i in range(int(num_images)):
-        # 経過時間をチェック
-        elapsed_time = time.time() - start_time
-        if elapsed_time >= duration - 5:  # duration変数を使用
-            break
-
         current_seed = seed + i if seed is not None and seed != -1 else None
         generator = (
             torch.Generator(device).manual_seed(int(current_seed)) if current_seed is not None else None
@@ -276,6 +263,7 @@ def start_tryon(
     # 最終的な結果を返す
     return output_images.value, mask_gray
 
+
 garm_list = os.listdir(os.path.join(example_path, "cloth"))
 garm_list_path = [os.path.join(example_path, "cloth", garm) for garm in garm_list]
 
@@ -341,7 +329,7 @@ with image_blocks as demo:
                 )
                 seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
                 num_images = gr.Slider(
-                    label="Number of Images", minimum=1, maximum=10, step=1, value=1
+                    label="Number of Images", minimum=1, maximum=10, step=1, value=1  # 最大値を10に変更
                 )
 
     try_button.click(