Update app.py

app.py

@@ -11,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
@@ -25,10 +25,19 @@ import cv2
 import os
 import requests
 
+@spaces.GPU
+def warmup_gpu():
+    import torch
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return True
+
+
 snapshot_download(repo_id="LiuZichen/MagicQuill-models", repo_type="model", local_dir="models")
 # HF_TOKEN = os.environ.get("HF_TOKEN")
 # The client has been made public. Welcome to duplicate our repo.
 client = Client("LiuZichen/DrawNGuess")
+warmup_gpu()
 scribbleColorEditModel = ScribbleColorEditModel()
 
 def tensor_to_numpy(tensor):
@@ -36,6 +45,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())
@@ -145,7 +210,7 @@ def guess_prompt_handler(original_image, add_color_image, add_edge_image):
 
     return res
     
-@spaces.GPU
+@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
@@ -178,7 +243,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):
@@ -330,4 +398,3 @@ app = gr.mount_gradio_app(app, demo, "/")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
-    # demo.launch()