Update app.py

app.py

@@ -122,7 +122,7 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 
-@spaces.GPU(duration=120)  # 実行時間を120秒に設定
+@spaces.GPU(duration=110)  # 実行時間を110秒に設定
 def start_tryon(
     dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, num_images
 ):
@@ -142,7 +142,7 @@ def start_tryon(
         left = (width - target_width) / 2
         top = (height - target_height) / 2
         right = (width + target_width) / 2
-        bottom = (height + target_height) / 2
+        bottom = (height + target_width) / 2
         cropped_img = human_img_orig.crop((left, top, right, bottom))
         crop_size = cropped_img.size
         human_img = cropped_img.resize((768, 1024))
@@ -178,7 +178,42 @@ def start_tryon(
     pose_img = pose_img[:, :, ::-1]
     pose_img = Image.fromarray(pose_img).resize((768, 1024))
 
-    output_images = []
+    # テキストエンコーディングは一度だけ行う
+    with torch.no_grad():
+        with torch.cuda.amp.autocast():
+            prompt = "model is wearing " + garment_des
+            negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
+            (
+                prompt_embeds,
+                negative_prompt_embeds,
+                pooled_prompt_embeds,
+                negative_pooled_prompt_embeds,
+            ) = pipe.encode_prompt(
+                prompt,
+                num_images_per_prompt=1,
+                do_classifier_free_guidance=True,
+                negative_prompt=negative_prompt,
+            )
+
+            prompt_c = "a photo of " + garment_des
+            negative_prompt_c = "monochrome, lowres, bad anatomy, worst quality, low quality"
+            (
+                prompt_embeds_c,
+                _,
+                _,
+                _,
+            ) = pipe.encode_prompt(
+                prompt_c,
+                num_images_per_prompt=1,
+                do_classifier_free_guidance=False,
+                negative_prompt=negative_prompt_c,
+            )
+
+    pose_img_tensor = tensor_transfrom(pose_img).unsqueeze(0).to(device, torch.float16)
+    garm_tensor = tensor_transfrom(garm_img).unsqueeze(0).to(device, torch.float16)
+
+    # 最初にマスク画像を一度だけ出力
+    yield None, mask_gray
 
     for i in range(int(num_images)):
         current_seed = seed + i if seed is not None and seed != -1 else None
@@ -188,37 +223,6 @@ def start_tryon(
 
         with torch.no_grad():
             with torch.cuda.amp.autocast():
-                prompt = "model is wearing " + garment_des
-                negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
-                (
-                    prompt_embeds,
-                    negative_prompt_embeds,
-                    pooled_prompt_embeds,
-                    negative_pooled_prompt_embeds,
-                ) = pipe.encode_prompt(
-                    prompt,
-                    num_images_per_prompt=1,
-                    do_classifier_free_guidance=True,
-                    negative_prompt=negative_prompt,
-                )
-
-                prompt_c = "a photo of " + garment_des
-                negative_prompt_c = "monochrome, lowres, bad anatomy, worst quality, low quality"
-                (
-                    prompt_embeds_c,
-                    _,
-                    _,
-                    _,
-                ) = pipe.encode_prompt(
-                    prompt_c,
-                    num_images_per_prompt=1,
-                    do_classifier_free_guidance=False,
-                    negative_prompt=negative_prompt_c,
-                )
-
-                pose_img_tensor = tensor_transfrom(pose_img).unsqueeze(0).to(device, torch.float16)
-                garm_tensor = tensor_transfrom(garm_img).unsqueeze(0).to(device, torch.float16)
-
                 images = pipe(
                     prompt_embeds=prompt_embeds.to(device, torch.float16),
                     negative_prompt_embeds=negative_prompt_embeds.to(device, torch.float16),
@@ -244,11 +248,9 @@ def start_tryon(
             out_img = images[0].resize(crop_size)
             human_img_copy = human_img_orig.copy()
             human_img_copy.paste(out_img, (int(left), int(top)))
-            output_images.append(human_img_copy)
+            yield human_img_copy, None
         else:
-            output_images.append(images[0])
-
-    return output_images, mask_gray
+            yield images[0], None
 
 
 garm_list = os.listdir(os.path.join(example_path, "cloth"))
@@ -299,9 +301,11 @@ with image_blocks as demo:
                     )
             example = gr.Examples(inputs=garm_img, examples_per_page=8, examples=garm_list_path)
         with gr.Column():
-            masked_img = gr.Image(label="Masked image output", elem_id="masked-img", show_share_button=False)
+            masked_img = gr.Image(
+                label="Masked image output", elem_id="masked-img", show_share_button=False
+            )
         with gr.Column():
-            image_out = gr.Gallery(label="Output", elem_id="output-img", show_share_button=False)
+            image_out = gr.Image(label="Output", elem_id="output-img", show_share_button=False)
 
     with gr.Column():
         try_button = gr.Button(value="Try-on")