Update generate.py

generate.py

@@ -1,7 +1,7 @@
 # generate.py
-# --- VERSION 12 (Final Fix) ---
+# --- VERSION 13 (Correct Upscaler Loading) ---
 
-print("--- RUNNING GENERATE.PY VERSION 12 (Final Fix) ---")
+print("--- RUNNING GENERATE.PY VERSION 13 (Correct Upscaler Loading) ---")
 
 # --- MONKEY-PATCH FOR OLD TORCHVISION ---
 try:
@@ -26,9 +26,9 @@ 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 realesrgan import RealESRGANer # <-- IMPORT THE CORRECT CLASS
+from basicsr.archs.rrdbnet_arch import RRDBNet
 from gfpgan import GFPGANer
-from basicsr.utils.download_util import load_file_from_url
 
 import config
 import utils
@@ -70,19 +70,20 @@ class GenerationService:
                 vae=vae, feature_extractor=None, safety_checker=None
             ).to(self.device)
             
-            # --- Upscaler Model ---
+            # --- CORRECTED UPSCALER LOADING ---
             logger.info("Loading Real-ESRGAN upscaler model...")
-            # 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 = 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')
-            self.upsampler.load_state_dict(torch.load(model_path)['params_ema'])
-            self.upsampler.to(self.device)
+            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+            self.upsampler = RealESRGANer(
+                scale=4,
+                model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth',
+                dni_weight=None,
+                model=model,
+                tile=0,
+                tile_pad=10,
+                pre_pad=0,
+                half=True if self.torch_dtype == torch.float16 else False,
+                gpu_id=0 if self.device == "cuda" else None
+            )
             logger.info("Upscaler model loaded.")
 
             logger.info("All models loaded successfully.")
@@ -95,16 +96,9 @@ class GenerationService:
         """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')
-            
-            cv2.imwrite(image_path, output_img)
+            # The enhance method returns the upscaled image and its type
+            output, _ = self.upsampler.enhance(img, outscale=4)
+            cv2.imwrite(image_path, output)
             logger.info(f"Successfully upscaled image: {image_path}")
             return image_path
         except Exception as e: