Update app.py

app.py

@@ -217,48 +217,55 @@ def start_tryon(
     # 最初にマスク画像を一度だけ出力
     yield output_images.value, mask_gray
 
-    for i in range(int(num_images)):
-        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
-        )
+    try:
+        for i in range(int(num_images)):
+            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
+            )
 
-        with torch.no_grad():
-            with torch.cuda.amp.autocast():
-                images = pipe(
-                    prompt_embeds=prompt_embeds.to(device, torch.float16),
-                    negative_prompt_embeds=negative_prompt_embeds.to(device, torch.float16),
-                    pooled_prompt_embeds=pooled_prompt_embeds.to(device, torch.float16),
-                    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(
-                        device, torch.float16
-                    ),
-                    num_inference_steps=denoise_steps,
-                    generator=generator,
-                    strength=1.0,
-                    pose_img=pose_img_tensor.to(device, torch.float16),
-                    text_embeds_cloth=prompt_embeds_c.to(device, torch.float16),
-                    cloth=garm_tensor.to(device, torch.float16),
-                    mask_image=mask,
-                    image=human_img,
-                    height=1024,
-                    width=768,
-                    ip_adapter_image=garm_img.resize((768, 1024)),
-                    guidance_scale=2.0,
-                )[0]
-
-        if is_checked_crop:
-            out_img = images[0].resize(crop_size)
-            human_img_copy = human_img_orig.copy()
-            human_img_copy.paste(out_img, (int(left), int(top)))
-            new_image = human_img_copy
-        else:
-            new_image = images[0]
-
-        # 画像を追加
-        output_images.value.append(new_image)
-
-        # ギャラリーを更新しつつmasked_imgを保持
-        yield output_images.value, mask_gray  # mask_grayを再度渡す
+            with torch.no_grad():
+                with torch.cuda.amp.autocast():
+                    images = pipe(
+                        prompt_embeds=prompt_embeds.to(device, torch.float16),
+                        negative_prompt_embeds=negative_prompt_embeds.to(device, torch.float16),
+                        pooled_prompt_embeds=pooled_prompt_embeds.to(device, torch.float16),
+                        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(
+                            device, torch.float16
+                        ),
+                        num_inference_steps=denoise_steps,
+                        generator=generator,
+                        strength=1.0,
+                        pose_img=pose_img_tensor.to(device, torch.float16),
+                        text_embeds_cloth=prompt_embeds_c.to(device, torch.float16),
+                        cloth=garm_tensor.to(device, torch.float16),
+                        mask_image=mask,
+                        image=human_img,
+                        height=1024,
+                        width=768,
+                        ip_adapter_image=garm_img.resize((768, 1024)),
+                        guidance_scale=2.0,
+                    )[0]
+
+            if is_checked_crop:
+                out_img = images[0].resize(crop_size)
+                human_img_copy = human_img_orig.copy()
+                human_img_copy.paste(out_img, (int(left), int(top)))
+                new_image = human_img_copy
+            else:
+                new_image = images[0]
+
+            # 画像を追加
+            output_images.value.append(new_image)
+
+            # ギャラリーを更新しつつmasked_imgを保持
+            yield output_images.value, mask_gray  # mask_grayを再度渡す
+
+    except Exception as e:
+        # タイムアウトなどの例外が発生した場合でも、生成済みの画像を返却
+        yield output_images.value, mask_gray
+        # エラーを再度発生させてエラーメッセージを表示
+        raise e
 
     # 最終的な結果を返す
     return output_images.value, mask_gray
@@ -316,7 +323,9 @@ with image_blocks as demo:
                 label="Masked image output", elem_id="masked-img", show_share_button=False
             )
         with gr.Column():
-            image_gallery = gr.Gallery(label="Generated Images", elem_id="output-gallery", show_label=True)
+            image_gallery = gr.Gallery(
+                label="Generated Images", elem_id="output-gallery", show_label=True
+            )
 
     with gr.Column():
         try_button = gr.Button(value="Try-on")
@@ -327,7 +336,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(
@@ -346,4 +355,4 @@ with image_blocks as demo:
         api_name='tryon',
     )
 
-image_blocks.launch(show_error=True)
+image_blocks.launch(show_error=True)