Update app.py

app.py

@@ -3,6 +3,8 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
+from PIL import Image, ImageDraw, ImageFont
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -18,26 +20,6 @@ else:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 3072
 
-# def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-        
-#     generator = torch.Generator().manual_seed(seed)
-    
-#     image = pipe(
-#         prompt = prompt, 
-#         negative_prompt = negative_prompt,
-#         guidance_scale = guidance_scale, 
-#         num_inference_steps = num_inference_steps, 
-#         width = width, 
-#         height = height,
-#         generator = generator
-#     ).images[0] 
-    
-#     return image
-
-
 
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     # Append Moroccan and Amazigh art styles to the prompt
@@ -45,20 +27,35 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
+
     generator = torch.Generator().manual_seed(seed)
-    
+
     image = pipe(
-        prompt = style_prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image
+        prompt=style_prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator
+    ).images[0]
+
+    # Convert the image to PIL format for overlaying the watermark
+    pil_image = Image.fromarray(image)
+
+    # Add watermark
+    watermark_text = "Bibou.jpeg"
+    font = ImageFont.truetype("arial.ttf", size=30)  # Adjust font and size as needed
+    draw = ImageDraw.Draw(pil_image)
+    text_width, text_height = draw.textsize(watermark_text, font=font)
+    margin = 10
+    opacity = 0.6
+    draw.text((pil_image.width - text_width - margin, pil_image.height - text_height - margin), watermark_text, font=font, fill=(255, 255, 255, int(255 * opacity)))
+
+    # Convert back to numpy array for Gradio display
+    watermarked_image = np.array(pil_image)
+
+    return watermarked_image