update: code refactor

app.py

@@ -32,8 +32,8 @@ class NecklaceTryOn:
         torch.cuda.empty_cache()
         gc.collect()
 
-    @spaces.GPU
-    def clothing_try_on_n_necklace_try_on(self, image, jewellery):
+    def apply_necklace(self, image, jewellery):
+        """Apply necklace on the image and return modified image and mask."""
         image = np.array(image)
         copy_image = image.copy()
         jewellery = np.array(jewellery)
@@ -97,26 +97,30 @@ class NecklaceTryOn:
 
         gc.collect()
 
-        image = Image.fromarray(result.astype(np.uint8))
-        mask = Image.fromarray(binaryMask.astype(np.uint8)).convert("RGB")
+        return Image.fromarray(result.astype(np.uint8)), Image.fromarray(binaryMask.astype(np.uint8)).convert("RGB")
+
+    @spaces.GPU
+    def clothing_try_on_n_necklace_try_on(self, image, jewellery):
+        """Main method for clothing and necklace try-on."""
+        result_image, mask = self.apply_necklace(image, jewellery)
 
         jewellery_mask = Image.fromarray(
-            np.bitwise_and(np.array(mask), np.array(image))
+            np.bitwise_and(np.array(mask), np.array(result_image))
         )
         arr_orig = np.array(grayscale(mask))
 
-        image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
-        image = Image.fromarray(image)
+        result_image = cv2.inpaint(np.array(result_image), arr_orig, 15, cv2.INPAINT_TELEA)
+        result_image = Image.fromarray(result_image)
 
         arr = arr_orig.copy()
         mask_y = np.where(arr == arr[arr != 0][0])[0][0]
         arr[mask_y:, :] = 255
 
-        new = Image.fromarray(arr)
-        mask = new.copy()
+        new_mask = Image.fromarray(arr)
+        mask = new_mask.copy()
 
-        orig_size = image.size
-        image = image.resize((512, 512))
+        orig_size = result_image.size
+        result_image = result_image.resize((512, 512))
         mask = mask.resize((512, 512))
 
         results = []
@@ -126,7 +130,7 @@ class NecklaceTryOn:
         output = self.pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            image=image,
+            image=result_image,
             mask_image=mask,
             strength=0.95,
             guidance_score=9,