Update generate.py

generate.py

@@ -1,7 +1,7 @@
 # generate.py
-# --- VERSION 8 (FINAL) ---
+# --- VERSION 9 (with Watermarking) ---
 
-print("--- RUNNING GENERATE.PY VERSION 8 (FINAL) ---")
+print("--- RUNNING GENERATE.PY VERSION 9 (with Watermarking) ---")
 
 import torch
 import cv2
@@ -17,6 +17,7 @@ from insightface.utils import face_align
 from huggingface_hub import hf_hub_download
 from storage3.utils import StorageException
 import config
+import utils # <-- IMPORT THE NEW UTILS FILE
 from database import supabase
 
 # --- Setup Logging ---
@@ -61,7 +62,7 @@ class GenerationService:
             raise RuntimeError(f"Could not initialize GenerationService: {e}") from e
 
     def generate_magic_image(self, face_images: list, gender: str, prompt: str, plan: str = 'free') -> str | None:
-        logger.info("Starting image generation process...")
+        logger.info(f"Starting image generation process for a user on the '{plan}' plan.")
         
         full_prompt = f"{prompt}, 4k, high-resolution, photorealistic, masterpiece, single person, solo portrait, centered composition"
         negative_prompt = "multiple people, group photo, crowd, two faces, three faces, multiple faces, collage, ugly, deformed, blurry, low quality"
@@ -88,13 +89,8 @@ class GenerationService:
 
         logger.info("Calling the generation pipeline...")
         try:
-            # --- FINAL FIX: Create negative embeddings ---
-            # The IP-Adapter expects both positive and negative image embeddings.
-            # We create a zeroed-out tensor for the negative part.
             positive_embedding = average_embedding.unsqueeze(0)
             negative_embedding = torch.zeros_like(positive_embedding)
-            
-            # Combine them for the pipeline
             final_embedding = torch.cat([negative_embedding, positive_embedding], dim=0)
 
             output = self.pipe(
@@ -117,6 +113,11 @@ class GenerationService:
             local_path = os.path.join(temp_dir, f"{uuid.uuid4()}.png")
             image.save(local_path)
             
+            # --- APPLY WATERMARK FOR FREE USERS ---
+            if plan == 'free':
+                utils.add_watermark(local_path, "@MagicFaceBot")
+            
+            # --- Upload to Supabase Storage ---
             storage_path = f"public/{os.path.basename(local_path)}"
             logger.info(f"Uploading {local_path} to Supabase bucket '{config.SUPABASE_BUCKET_NAME}' at path '{storage_path}'")
             
@@ -150,10 +151,12 @@ if __name__ == '__main__':
         result_url = service.generate_magic_image(
             face_images=["test_face.jpg"],
             gender="Female",
-            prompt="A beautiful portrait of a princess in a magical forest, fantasy art"
+            prompt="A beautiful portrait of a princess in a magical forest, fantasy art",
+            plan='free' # <-- Test the 'free' plan to see the watermark
         )
         if result_url:
             print(f"\n✅ Test successful! Image URL: {result_url}")
+            print("Check the image at the URL to see if it has a watermark.")
         else:
             print(f"\n❌ Test failed. Please check the new traceback logs for details.")
     else: