Update app.py

app.py

@@ -178,81 +178,77 @@ def start_tryon(
     pose_img = pose_img[:, :, ::-1]
     pose_img = Image.fromarray(pose_img).resize((768, 1024))
 
-    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=num_images,
-                do_classifier_free_guidance=True,
-                negative_prompt=negative_prompt,
-            )
+    output_images = []
 
-            prompt = "a photo of " + garment_des
-            negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
-            if not isinstance(prompt, List):
-                prompt = [prompt] * num_images
-            if not isinstance(negative_prompt, List):
-                negative_prompt = [negative_prompt] * num_images
-            (
-                prompt_embeds_c,
-                _,
-                _,
-                _,
-            ) = pipe.encode_prompt(
-                prompt,
-                num_images_per_prompt=num_images,
-                do_classifier_free_guidance=False,
-                negative_prompt=negative_prompt,
-            )
+    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():
+                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,
+                )
 
-            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)
-            pose_img_tensor = pose_img_tensor.repeat(num_images, 1, 1, 1)
-            garm_tensor = garm_tensor.repeat(num_images, 1, 1, 1)
-            human_imgs = [human_img] * num_images
-            masks = [mask] * num_images
-            ip_adapter_images = [garm_img.resize((768, 1024))] * num_images
-
-            if seed is not None and seed != -1:
-                generator = [torch.Generator(device).manual_seed(seed + i) for i in range(num_images)]
-            else:
-                generator = None
-
-            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=masks,
-                image=human_imgs,
-                height=1024,
-                width=768,
-                ip_adapter_image=ip_adapter_images,
-                guidance_scale=2.0,
-            )[0]
+                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,
+                )
 
-    if is_checked_crop:
-        output_images = []
-        for img in images:
-            out_img = img.resize(crop_size)
-            human_img_orig.paste(out_img, (int(left), int(top)))
-            output_images.append(human_img_orig.copy())
-        return output_images, mask_gray
-    else:
-        return images, mask_gray
+                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),
+                    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)))
+            output_images.append(human_img_copy)
+        else:
+            output_images.append(images[0])
+
+    return output_images, mask_gray
 
 
 garm_list = os.listdir(os.path.join(example_path, "cloth"))