Update generate.py

generate.py

@@ -1,12 +1,9 @@
 # generate.py
-# --- VERSION 11 (with Monkey-Patch) ---
+# --- VERSION 12 (Final Fix) ---
 
-print("--- RUNNING GENERATE.PY VERSION 11 (with Monkey-Patch) ---")
+print("--- RUNNING GENERATE.PY VERSION 12 (Final Fix) ---")
 
 # --- MONKEY-PATCH FOR OLD TORCHVISION ---
-# This block fixes the "ModuleNotFoundError" by tricking the system.
-# It tells any library looking for the new `functional_tensor` module
-# to use the old `functional` module instead.
 try:
     import sys
     import torchvision.transforms.functional as F
@@ -75,9 +72,12 @@ class GenerationService:
             
             # --- Upscaler Model ---
             logger.info("Loading Real-ESRGAN upscaler model...")
-            model_path = os.path.join('weights', 'realesrgan-x4plus.pth')
+            # FIX: The downloaded file has uppercase letters. Match the filename exactly.
+            model_name = 'RealESRGAN_x4plus.pth'
+            model_path = os.path.join('weights', model_name)
+            
             if not os.path.exists(model_path):
-                model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
+                model_url = f'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/{model_name}'
                 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')
@@ -104,13 +104,12 @@ class GenerationService:
             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
+            return image_path
 
 
     def generate_magic_image(self, face_images: list, gender: str, prompt: str, plan: str = 'free') -> str | None:
@@ -156,14 +155,11 @@ class GenerationService:
             local_path = os.path.join(temp_dir, f"{uuid.uuid4()}.png")
             image.save(local_path)
             
-            # --- 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)}"
             with open(local_path, 'rb') as f:
                 supabase.storage.from_(config.SUPABASE_BUCKET_NAME).upload(
@@ -191,7 +187,7 @@ if __name__ == '__main__':
             face_images=["test_face.jpg"],
             gender="Female",
             prompt="A beautiful portrait of a princess in a magical forest, fantasy art",
-            plan='paid' # <-- Test the 'paid' plan to see the upscaling
+            plan='paid'
         )
         if result_url:
             print(f"\n✅ Test successful! Image URL: {result_url}")