Update app.py

app.py

@@ -3,7 +3,7 @@ import spaces
 import torch
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import os
 
 # Style dictionary
@@ -88,7 +88,52 @@ thumbnail_mapping = {
 pipeline = None
 pipeline_loaded = False
 
-def load_pipeline():
+def add_watermark(image):
+    """Add watermark to the generated image"""
+    # Create a copy of the image
+    watermarked = image.copy()
+    
+    # Create drawing context
+    draw = ImageDraw.Draw(watermarked)
+    
+    # Define watermark text
+    watermark_text = "ginigen.com"
+    
+    # Try to use a default font, fallback to basic if not available
+    try:
+        # Try to use a smaller font size
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16)
+    except:
+        try:
+            font = ImageFont.truetype("arial.ttf", 16)
+        except:
+            # Use default font if no TrueType fonts are available
+            font = ImageFont.load_default()
+    
+    # Get text size
+    bbox = draw.textbbox((0, 0), watermark_text, font=font)
+    text_width = bbox[2] - bbox[0]
+    text_height = bbox[3] - bbox[1]
+    
+    # Calculate position (bottom right with padding)
+    padding = 10
+    box_padding = 5
+    x = image.width - text_width - padding - (box_padding * 2)
+    y = image.height - text_height - padding - (box_padding * 2)
+    
+    # Draw semi-transparent box background
+    box_coords = [
+        x - box_padding,
+        y - box_padding,
+        x + text_width + box_padding,
+        y + text_height + box_padding
+    ]
+    draw.rectangle(box_coords, fill=(255, 255, 255, 200))
+    
+    # Draw text
+    draw.text((x, y), watermark_text, font=font, fill=(100, 100, 100, 255))
+    
+    return watermarked
     global pipeline, pipeline_loaded
     if pipeline is None:
         print("Loading FLUX.1-Kontext-dev model...")
@@ -182,7 +227,11 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         # Clear GPU memory
         torch.cuda.empty_cache()
         
-        return result.images[0]
+        # Add watermark to the generated image
+        result_image = result.images[0]
+        watermarked_image = add_watermark(result_image)
+        
+        return watermarked_image
         
     except Exception as e:
         print(f"Error: {str(e)}")
@@ -215,7 +264,7 @@ def create_thumbnail_grid():
 
 # Create Gradio interface
 with gr.Blocks(title="Flux Kontext Style LoRA", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 Flux Styler : Flux Kontext Style LoRA")
+    gr.Markdown("# 🎨 Flux Kontext Style LoRA")
     
     # Thumbnail Grid Section
     with gr.Row():