Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import subprocess
 import shlex
 subprocess.run(
@@ -10,7 +11,7 @@ from gradio_magicquill import MagicQuill
 import random
 import torch
 import numpy as np
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageDraw, ImageFont
 import base64
 import io
 from fastapi import FastAPI, Request
@@ -35,6 +36,62 @@ def tensor_to_numpy(tensor):
         return (tensor.detach().cpu().numpy() * 255).astype(np.uint8)
     return tensor
 
+def add_watermark_to_image(image_tensor, watermark_text="Power By magicquill.online"):
+    """
+    在图像右下角添加文字水印
+    """
+    # 将tensor转换为PIL图像
+    if isinstance(image_tensor, torch.Tensor):
+        image_array = (image_tensor.squeeze(0).detach().cpu().numpy() * 255).astype(np.uint8)
+    else:
+        image_array = image_tensor
+    
+    pil_image = Image.fromarray(image_array)
+    
+    # 获取图像尺寸
+    width, height = pil_image.size
+    
+    # 尝试加载字体，如果失败则使用默认字体
+    try:
+        # 根据图像大小动态调整字体大小
+        font_size = max(12, min(width, height) // 50)
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", font_size)
+    except:
+        try:
+            font_size = max(12, min(width, height) // 50)
+            font = ImageFont.load_default()
+        except:
+            font = ImageFont.load_default()
+    
+    # 创建绘制对象
+    draw = ImageDraw.Draw(pil_image)
+    
+    # 获取文字尺寸
+    bbox = draw.textbbox((0, 0), watermark_text, font=font)
+    text_width = bbox[2] - bbox[0]
+    text_height = bbox[3] - bbox[1]
+    
+    # 计算水印位置（右下角，留一些边距）
+    margin = 10
+    background_padding = 3  # 减小背景padding
+    x = width - text_width - margin - background_padding
+    y = height - text_height - margin - background_padding
+    
+    # 绘制紧贴文字的半透明背景矩形
+    draw.rectangle([
+        x - background_padding, 
+        y - background_padding,
+        x + text_width + background_padding,
+        y + text_height + background_padding
+    ], fill=(0, 0, 0, 160))  # 半透明黑色背景，稍微增加透明度
+    
+    # 绘制白色文字
+    draw.text((x, y), watermark_text, fill=(255, 255, 255, 255), font=font)
+    
+    # 转换回tensor格式
+    image_array = np.array(pil_image).astype(np.float32) / 255.0
+    return torch.from_numpy(image_array).unsqueeze(0)
+
 def tensor_to_base64(tensor):
     tensor = tensor.squeeze(0) * 255.
     pil_image = Image.fromarray(tensor.cpu().byte().numpy())
@@ -143,7 +200,8 @@ def guess_prompt_handler(original_image, add_color_image, add_edge_image):
         os.remove(add_edge_mask_file.name)
 
     return res
-
+    
+@spaces.GPU(duration=120) 
 def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_image, remove_edge_image, positive_prompt, negative_prompt, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
     add_color_image, original_image, total_mask, add_edge_mask, remove_edge_mask = prepare_images_and_masks(total_mask, original_image, add_color_image, add_edge_image, remove_edge_image)
     progress = None
@@ -176,7 +234,10 @@ def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_im
         progress
     )
 
-    final_image_base64 = tensor_to_base64(final_image)
+    # 在最终图像上添加水印
+    final_image_with_watermark = add_watermark_to_image(final_image, "Power By magicquill.online")
+    
+    final_image_base64 = tensor_to_base64(final_image_with_watermark)
     return final_image_base64
 
 def generate_image_handler(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
@@ -196,19 +257,15 @@ css = '''
 }
 '''
 
-head = """
-<meta http-equiv="Content-Security-Policy" content="frame-ancestors 'none'">
-"""
-
-with gr.Blocks(css=css, head=head) as demo:
+with gr.Blocks(css=css) as demo:
     with gr.Row(elem_classes="row"):
         text = gr.Markdown(
         """
-        # Welcome to MagicQuill! The paper has been accepted to CVPR 2025.
-        Click the [link](https://magicquill.art) to view our demo and tutorial. Give us a [GitHub star](https://github.com/magic-quill/magicquill) if you are interested. 
+        # Welcome to MagicQuill!
+        Thank you to the developers for their contributions. Give a [GitHub star](https://github.com/magic-quill/magicquill)
         """)
     with gr.Row(elem_classes="row"):
-        ms = MagicQuill()
+        ms = MagicQuill(theme="light")
     with gr.Row(elem_classes="row"):
         with gr.Column():
             btn = gr.Button("Run", variant="primary")
@@ -325,10 +382,10 @@ async def process_background_img(request: Request):
     img = await request.json()
     resized_img_tensor = load_and_resize_image(img)
     resized_img_base64 = "data:image/png;base64," + tensor_to_base64(resized_img_tensor)
+    # add more processing here
     return resized_img_base64
 
 app = gr.mount_gradio_app(app, demo, "/")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
-    # demo.launch()