Update generate.py

generate.py

@@ -1,7 +1,7 @@
 # generate.py
-# --- VERSION 9 (with Watermarking) ---
+# --- VERSION 10 (with Upscaling) ---
 
-print("--- RUNNING GENERATE.PY VERSION 9 (with Watermarking) ---")
+print("--- RUNNING GENERATE.PY VERSION 10 (with Upscaling) ---")
 
 import torch
 import cv2
@@ -16,8 +16,12 @@ from insightface.app import FaceAnalysis
 from insightface.utils import face_align
 from huggingface_hub import hf_hub_download
 from storage3.utils import StorageException
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan import GFPGANer
+from basicsr.utils.download_util import load_file_from_url
+
 import config
-import utils # <-- IMPORT THE NEW UTILS FILE
+import utils
 from database import supabase
 
 # --- Setup Logging ---
@@ -38,6 +42,7 @@ class GenerationService:
         vae_model_path = "stabilityai/sd-vae-ft-mse"
         
         try:
+            # --- AI Models ---
             self.face_app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider' if self.device == "cuda" else 'CPUExecutionProvider'])
             self.face_app.prepare(ctx_id=0, det_size=(640, 640))
             cv2.setNumThreads(1)
@@ -55,12 +60,46 @@ class GenerationService:
                 vae=vae, feature_extractor=None, safety_checker=None
             ).to(self.device)
             
+            # --- Upscaler Model ---
+            logger.info("Loading Real-ESRGAN upscaler model...")
+            model_path = os.path.join('weights', 'realesrgan-x4plus.pth')
+            if not os.path.exists(model_path):
+                model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
+                load_file_from_url(url=model_url, model_dir=os.path.join('weights'), progress=True)
+            
+            self.upsampler = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+            self.upsampler.load_state_dict(torch.load(model_path)['params_ema'])
+            self.upsampler.to(self.device)
+            logger.info("Upscaler model loaded.")
+
             logger.info("All models loaded successfully.")
 
         except Exception as e:
             logger.error(f"Fatal error during model loading: {e}")
             raise RuntimeError(f"Could not initialize GenerationService: {e}") from e
 
+    def _upscale_image(self, image_path: str) -> str:
+        """Upscales an image using Real-ESRGAN."""
+        try:
+            img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
+            img = img.astype('float32') / 255.
+            img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(self.device)
+
+            with torch.no_grad():
+                output = self.upsampler(img)
+
+            output_img = output.squeeze().permute(1, 2, 0).cpu().numpy()
+            output_img = (output_img * 255.0).round().astype('uint8')
+            
+            # Save the upscaled image back to the same path
+            cv2.imwrite(image_path, output_img)
+            logger.info(f"Successfully upscaled image: {image_path}")
+            return image_path
+        except Exception as e:
+            logger.error(f"Failed to upscale image {image_path}: {e}")
+            return image_path # Return original path on failure
+
+
     def generate_magic_image(self, face_images: list, gender: str, prompt: str, plan: str = 'free') -> str | None:
         logger.info(f"Starting image generation process for a user on the '{plan}' plan.")
         
@@ -68,12 +107,10 @@ class GenerationService:
         negative_prompt = "multiple people, group photo, crowd, two faces, three faces, multiple faces, collage, ugly, deformed, blurry, low quality"
 
         faceid_all_embeds = []
-        
         for image_path in face_images:
             try:
                 face = cv2.imread(image_path)
                 if face is None: continue
-                
                 faces = self.face_app.get(face)
                 if faces:
                     faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
@@ -94,43 +131,32 @@ class GenerationService:
             final_embedding = torch.cat([negative_embedding, positive_embedding], dim=0)
 
             output = self.pipe(
-                prompt=full_prompt,
-                negative_prompt=negative_prompt,
-                ip_adapter_image_embeds=[final_embedding],
-                num_inference_steps=40,
-                guidance_scale=7.5,
-                width=512,
-                height=768,
+                prompt=full_prompt, negative_prompt=negative_prompt,
+                ip_adapter_image_embeds=[final_embedding], num_inference_steps=40,
+                guidance_scale=7.5, width=512, height=768,
             )
             
-            if isinstance(output, StableDiffusionPipelineOutput):
-                image = output.images[0]
-            else:
-                image = output[0][0]
+            image = output.images[0] if isinstance(output, StableDiffusionPipelineOutput) else output[0][0]
 
             temp_dir = "temp_images"
             os.makedirs(temp_dir, exist_ok=True)
             local_path = os.path.join(temp_dir, f"{uuid.uuid4()}.png")
             image.save(local_path)
             
-            # --- APPLY WATERMARK FOR FREE USERS ---
+            # --- FEATURE TIER LOGIC ---
             if plan == 'free':
                 utils.add_watermark(local_path, "@MagicFaceBot")
+            else:
+                # Upscale for paid users
+                self._upscale_image(local_path)
             
             # --- 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}'")
-            
             with open(local_path, 'rb') as f:
                 supabase.storage.from_(config.SUPABASE_BUCKET_NAME).upload(
-                    path=storage_path,
-                    file=f,
-                    file_options={"content-type": "image/png"}
+                    path=storage_path, file=f, file_options={"content-type": "image/png"}
                 )
-            
             public_url = supabase.storage.from_(config.SUPABASE_BUCKET_NAME).get_public_url(storage_path)
-            logger.info(f"Upload successful. Public URL: {public_url}")
-
             os.remove(local_path)
             
             return public_url
@@ -152,12 +178,12 @@ if __name__ == '__main__':
             face_images=["test_face.jpg"],
             gender="Female",
             prompt="A beautiful portrait of a princess in a magical forest, fantasy art",
-            plan='free' # <-- Test the 'free' plan to see the watermark
+            plan='paid' # <-- Test the 'paid' plan to see the upscaling
         )
         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.")
+            print("Check the image at the URL. It should be high-resolution and have no watermark.")
         else:
-            print(f"\n❌ Test failed. Please check the new traceback logs for details.")
+            print(f"\n❌ Test failed. Please check the logs for details.")
     else:
         print("To run a test, place an image named 'test_face.jpg' in the root directory.")