Update main_code_script.py

main_code_script.py

@@ -54,21 +54,20 @@ def segment_clothing(image, results): #Added result
     mask_img = Image.fromarray(binary_mask).convert("L")
     return mask_img
 # --- 3. Image Inpainting (Replacing Clothing - using Stable Diffusion Inpainting) ---
-def inpaint_clothing(image, mask_img, clothing_prompt, device="cuda" if torch.cuda.is_available() else "cpu"):
+def inpaint_clothing(image, mask_img, clothing_prompt, device="cpu" if torch.cuda.is_available() else "cpu"):
     """
     Replaces the clothing region in the image with new clothing based on a text prompt,
     using Stable Diffusion Inpainting.
     """
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
-#        "stabilityai/stable-diffusion-2-inpainting",
-        "stabilityai/stable-diffusion-2-base"
-#        "SimianLuo/LCM_Dreamshaper_v7"
+         "stabilityai/stable-diffusion-2-inpainting",
+         "SimianLuo/LCM_Dreamshaper_v7"
         torch_dtype=torch.float16
     )
     pipe = pipe.to(device)
     # Resize the image and mask to a smaller size for faster inpainting
     width, height = image.size
-    inpainted_size = (256, 256)  # Smaller size for faster inpainting
+    inpainted_size = (128, 128)  # Smaller size for faster inpainting
     image = image.resize(inpainted_size)
     mask_img = mask_img.resize(inpainted_size)
     prompt = f"A photo of a person wearing {clothing_prompt}" #Add style or detail