Update app.py

app.py

@@ -15,8 +15,52 @@ torch.jit.script = lambda f: f
 from model.cloth_masker import AutoMasker, vis_mask
 from model.pipeline import CatVTONPipeline
 from utils import init_weight_dtype, resize_and_crop, resize_and_padding
+from PIL import Image
 
-
+def add_watermark(main_image, logo_path, position='bottom-right', size_percentage=10):
+    """
+    Add a watermark to an image
+    Args:
+        main_image (PIL.Image): The main image
+        logo_path (str): Path to the logo image
+        position (str): Position of watermark ('bottom-right', 'bottom-left', 'top-right', 'top-left')
+        size_percentage (int): Size of watermark relative to main image (in percentage)
+    Returns:
+        PIL.Image: Image with watermark
+    """
+    # Open and resize the logo
+    logo = Image.open(logo_path).convert('RGBA')
+    
+    # Calculate the size for the logo
+    main_width, main_height = main_image.size
+    logo_width = int(main_width * size_percentage / 100)
+    logo_height = int(logo.size[1] * (logo_width / logo.size[0]))
+    logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
+    
+    # Convert main image to RGBA if it isn't already
+    if main_image.mode != 'RGBA':
+        main_image = main_image.convert('RGBA')
+    
+    # Create a new blank image with the same size as main image
+    watermarked = Image.new('RGBA', main_image.size, (0, 0, 0, 0))
+    watermarked.paste(main_image, (0, 0))
+    
+    # Calculate position
+    if position == 'bottom-right':
+        position = (main_width - logo_width - 10, main_height - logo_height - 10)
+    elif position == 'bottom-left':
+        position = (10, main_height - logo_height - 10)
+    elif position == 'top-right':
+        position = (main_width - logo_width - 10, 10)
+    elif position == 'top-left':
+        position = (10, 10)
+    
+    # Paste the logo
+    watermarked.paste(logo, position, logo)
+    
+    # Convert back to RGB
+    return watermarked.convert('RGB')
+    
 def parse_args():
     parser = argparse.ArgumentParser(description="Simple example of a training script.")
     parser.add_argument(
@@ -186,8 +230,9 @@ def submit_function(
     masked_person = vis_mask(person_image, mask)
     save_result_image = image_grid([person_image, masked_person, cloth_image, result_image], 1, 4)
     save_result_image.save(result_save_path)
+    final_image = None
     if show_type == "result only":
-        return result_image
+        final_image = result_image
     else:
         width, height = person_image.size
         if show_type == "input & result":
@@ -200,7 +245,11 @@ def submit_function(
         new_result_image = Image.new("RGB", (width + condition_width + 5, height))
         new_result_image.paste(conditions, (0, 0))
         new_result_image.paste(result_image, (condition_width + 5, 0))
-    return new_result_image
+        final_image = new_result_image
+    
+    # Add watermark
+    watermarked_image = add_watermark(final_image, 'logo.png', 'bottom-right', 5)
+    return watermarked_image
 
 
 def person_example_fn(image_path):