Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+mport gradio as gr
 import numpy as np
 import random
 import spaces
@@ -6,86 +6,96 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 import io
+from PIL import ImageEnhance
 
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-pipe = DiffusionPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell",
-    torch_dtype=dtype
-).to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
-
-import numpy as np
-from collections import Counter
-
-def get_prominent_colors(image, num_colors=5):
+def get_edge_color(image):
     """
-    Get the most prominent colors from an image, focusing on edges
+    Get a random color from the edge of the image
     """
     # Convert to numpy array
     img_array = np.array(image)
     
-    # Create a simple edge mask using gradient magnitude
-    gradient_x = np.gradient(img_array.mean(axis=2))[1]
-    gradient_y = np.gradient(img_array.mean(axis=2))[0]
-    gradient_magnitude = np.sqrt(gradient_x**2 + gradient_y**2)
+    # Get pixels from all edges
+    top_edge = img_array[0, :, :]
+    bottom_edge = img_array[-1, :, :]
+    left_edge = img_array[:, 0, :]
+    right_edge = img_array[:, -1, :]
+    
+    # Combine all edge pixels
+    edge_pixels = np.concatenate([top_edge, bottom_edge, left_edge, right_edge])
+    
+    # Pick a random edge pixel
+    random_edge_color = tuple(edge_pixels[random.randint(0, len(edge_pixels)-1)])
+    
+    return random_edge_color
+
+def color_match_tshirt(tshirt_image, target_color, threshold=30):
+    """
+    Change white/near-white areas of the t-shirt to the target color
+    """
+    # Convert to numpy array
+    img_array = np.array(tshirt_image)
+    
+    # Create a mask for near-white pixels
+    white_mask = np.all(np.abs(img_array - [255, 255, 255]) < threshold, axis=2)
+    
+    # Apply the new color to masked areas
+    img_array[white_mask] = target_color
     
-    # Threshold to get edge pixels
-    edge_threshold = np.percentile(gradient_magnitude, 90)  # Adjust percentile as needed
-    edge_mask = gradient_magnitude > edge_threshold
+    return Image.fromarray(img_array)
+
+def add_logo_watermark(image, logo_path='logo.png', size_percentage=0.2):
+    """
+    Add a logo watermark to the bottom right corner
+    """
+    # Open and resize logo
+    logo = Image.open(logo_path)
     
-    # Get colors from edge pixels
-    edge_colors = img_array[edge_mask]
+    # Calculate new logo size (20% of main image width)
+    new_width = int(image.size[0] * size_percentage)
+    new_height = int(new_width * logo.size[1] / logo.size[0])
+    logo = logo.resize((new_width, new_height), Image.Resampling.LANCZOS)
     
-    # Convert colors to tuples for counting
-    colors = [tuple(color) for color in edge_colors]
+    # If logo has transparency, use it as mask
+    if logo.mode == 'RGBA':
+        mask = logo.split()[3]
+    else:
+        mask = None
     
-    # Count occurrences of each color
-    color_counts = Counter(colors)
+    # Calculate position (bottom right corner with padding)
+    position = (image.size[0] - logo.size[0] - 10, 
+               image.size[1] - logo.size[1] - 10)
     
-    # Get most common colors
-    prominent_colors = color_counts.most_common(num_colors)
+    # Create a copy of the image
+    result = image.copy()
+    result.paste(logo, position, mask)
     
-    return prominent_colors
+    return result
 
 def create_tshirt_preview(design_image, tshirt_color="white"):
     """
-    Overlay the design onto the existing t-shirt template and color match
+    Overlay the design onto the t-shirt template with color matching
     """
     # Load the template t-shirt image
     tshirt = Image.open('image.jpeg')
     tshirt_width, tshirt_height = tshirt.size
     
-    # Convert design to PIL Image if it's not already
-    if not isinstance(design_image, Image.Image):
-        design_image = Image.fromarray(design_image)
-    
-    # Get prominent colors from the design
-    prominent_colors = get_prominent_colors(design_image)
-    if prominent_colors:
-        # Use the most prominent color for the t-shirt
-        main_color = prominent_colors[0][0]  # RGB tuple of most common color
-    else:
-        # Fallback to white if no colors found
-        main_color = (255, 255, 255)
+    # Get a random edge color from the design
+    edge_color = get_edge_color(design_image)
     
-    # Convert design to PIL Image if it's not already
-    if not isinstance(design_image, Image.Image):
-        design_image = Image.fromarray(design_image)
+    # Color match the t-shirt
+    tshirt = color_match_tshirt(tshirt, edge_color)
     
-    # Resize design to fit nicely on shirt (40% of shirt width)
-    design_width = int(tshirt_width * 0.35)  # Adjust this percentage as needed
+    # Resize design (35% of shirt width)
+    design_width = int(tshirt_width * 0.35)
     design_height = int(design_width * design_image.size[1] / design_image.size[0])
     design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
     
-    # Calculate position to center design on shirt
+    # Calculate position to center design
     x = (tshirt_width - design_width) // 2
-    y = int(tshirt_height * 0.2)  # Adjust this value based on your template
+    y = int(tshirt_height * 0.2)
     
-    # If design has transparency (RGBA), create mask
+    # Create mask if design has transparency
     if design_image.mode == 'RGBA':
         mask = design_image.split()[3]
     else:
@@ -94,34 +104,10 @@ def create_tshirt_preview(design_image, tshirt_color="white"):
     # Paste design onto shirt
     tshirt.paste(design_image, (x, y), mask)
     
-    return tshirt
-
-def enhance_prompt_for_tshirt(prompt, style=None):
-    """Add specific terms to ensure good t-shirt designs."""
-    style_terms = {
-        "minimal": ["simple geometric shapes", "clean lines", "minimalist illustration"],
-        "vintage": ["distressed effect", "retro typography", "vintage illustration"],
-        "artistic": ["hand-drawn style", "watercolor effect", "artistic illustration"],
-        "geometric": ["abstract shapes", "geometric patterns", "modern design"],
-        "typography": ["bold typography", "creative lettering", "text-based design"],
-        "realistic": ["realistic", "cinematic", "photograph"]
-    }
-    
-    base_terms = [
-        "create t-shirt design",
-        "with centered composition",
-        "high quality",
-        "professional design",
-        "clear background"
-    ]
-    
-    enhanced_prompt = f"{prompt}, {', '.join(base_terms)}"
+    # Add logo watermark
+    tshirt = add_logo_watermark(tshirt)
     
-    if style and style in style_terms:
-        style_specific_terms = style_terms[style]
-        enhanced_prompt = f"{enhanced_prompt}, {', '.join(style_specific_terms)}"
-    
-    return enhanced_prompt
+    return tshirt
 
 @spaces.GPU()
 def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=False, 
@@ -143,6 +129,9 @@ def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=Fals
         guidance_scale=0.0
     ).images[0]
     
+    # Add logo to design
+    design_image = add_logo_watermark(design_image)
+    
     # Create t-shirt preview
     tshirt_preview = create_tshirt_preview(design_image, tshirt_color)