Update inference_realesrgan_gpu.py

inference_realesrgan_gpu.py

@@ -3,9 +3,7 @@
 
 """
 GPU-only Real-ESRGAN + GFPGAN inference script.
-- Kötelezően CUDA-t követel (ha nincs GPU, hibát dob).
-- A modell(ek) és a GFPGAN belső hálói expliciten GPU-ra kerülnek.
-- Támogat fp16 (half) inference, ha nem adod meg --fp32.
+FIXED VERSION: Safe string handling & Forced FP32 for GFPGAN.
 """
 
 import argparse
@@ -24,25 +22,20 @@ from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
 
 def setup_device(gpu_id: int) -> torch.device:
-    """Ellenőrzi CUDA elérhetőségét, beállítja az eszközt és néhány gyorsítótár beállítást."""
+    """Ellenőrzi CUDA elérhetőségét, beállítja az eszközt."""
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA nem elérhető — ez a script csak GPU-n futtatható.")
-    # biztosítsuk, hogy a megadott GPU legyen kiválasztva
     torch.cuda.set_device(int(gpu_id))
     device = torch.device(f"cuda:{int(gpu_id)}")
     torch.backends.cudnn.benchmark = True
-    # kikapcsoljuk a gradet inferencia alatt
     torch.set_grad_enabled(False)
     return device
 
 
 def move_obj_to_device(obj: Any, device: torch.device, use_half: bool):
     """
-    Rekurzívan megpróbál minden torch.nn.Module objektumot GPU-ra mozgatni,
-    és a lehetőségek szerint half()-olni.
-    Ez a GFPGAN különböző verzióihoz hasznos.
+    Rekurzívan megpróbál minden torch.nn.Module objektumot GPU-ra mozgatni.
     """
-    # Modul esetén egyszerűen .to(device) és .half() ha lehet
     try:
         import torch.nn as nn
     except Exception:
@@ -55,12 +48,18 @@ def move_obj_to_device(obj: Any, device: torch.device, use_half: bool):
             pass
         if use_half:
             try:
+                # Csak akkor konvertáljuk, ha explicit kértük
                 obj.half()
             except Exception:
                 pass
+        else:
+            # Ha NEM kérünk half-ot, biztosítjuk, hogy float legyen (GFPGAN javítás)
+            try:
+                obj.float()
+            except Exception:
+                pass
         return
 
-    # ha dict/list/tuple, nézzük át az elemeket
     if isinstance(obj, dict):
         for v in obj.values():
             move_obj_to_device(v, device, use_half)
@@ -70,21 +69,16 @@ def move_obj_to_device(obj: Any, device: torch.device, use_half: bool):
             move_obj_to_device(v, device, use_half)
         return
 
-    # ha egy objektumnak vannak attribútumai, próbáljuk átnézni őket (GFPGAN belsők)
     if hasattr(obj, "__dict__"):
         for _, v in vars(obj).items():
-            # elkerüljük a végtelen rekurrenciát és az egyszerű típusokat
             if v is None:
                 continue
-            # közvetlenül modulok és konténerek kezelése
             try:
-                # modulokra alapból ráhívjuk a mozgást
                 if nn is not None and isinstance(v, nn.Module):
                     move_obj_to_device(v, device, use_half)
                 elif isinstance(v, (list, tuple, dict, set)):
                     move_obj_to_device(v, device, use_half)
             except Exception:
-                # tűrjük a hibákat, mert GFPGAN belsők különbözőek lehetnek
                 pass
 
 
@@ -111,9 +105,10 @@ def main():
     args = parser.parse_args()
 
     device = setup_device(args.gpu_id)
+    # Ha fp32 flag nincs megadva, akkor use_half=True
     use_half = not args.fp32
 
-    # model kiválasztása
+    # --- MODEL KIVÁLASZTÁS ---
     args.model_name = args.model_name.split('.')[0]
     if args.model_name == 'RealESRGAN_x4plus':
         model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
@@ -145,7 +140,6 @@ def main():
     else:
         raise ValueError(f"Ismeretlen model_name: {args.model_name}")
 
-    # model path letöltése ha szükséges
     if args.model_path is not None:
         model_path = args.model_path
     else:
@@ -155,24 +149,21 @@ def main():
             for url in file_url:
                 model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
 
-    # dni weight (realesr-general esetén)
     dni_weight = None
     if args.model_name == 'realesr-general-x4v3' and args.denoise_strength != 1:
         wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
         model_path = [model_path, wdn_model_path]
         dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
 
-    # modell GPU-ra helyezése és precision beállítása
+    # --- RealESRGAN (Háttér) Modell betöltése ---
     model = model.to(device)
     model.eval()
     if use_half:
         try:
             model.half()
         except Exception:
-            # ha nem támogatja a half-ot, megy fp32-ben
             print("Figyelem: modell nem támogatta a .half() hívást -> használ fp32-t.")
 
-    # RealESRGANer létrehozása (gpu_id explicit, hogy ne essen CPU fallbackbe)
     upsampler = RealESRGANer(
         scale=netscale,
         model_path=model_path,
@@ -182,10 +173,10 @@ def main():
         tile_pad=args.tile_pad,
         pre_pad=args.pre_pad,
         half=use_half,
-        gpu_id=int(args.gpu_id)  # soha ne legyen None
+        gpu_id=int(args.gpu_id)
     )
 
-    # GFPGAN inicializálása (ha kértük) és kényszerített GPU-ra mozgatás
+    # --- GFPGAN (Arc) Inicializálása ---
     face_enhancer = None
     if args.face_enhance:
         try:
@@ -193,48 +184,33 @@ def main():
         except Exception as e:
             raise RuntimeError("GFPGAN kértél, de a gfpgan modul nem található: " + str(e))
 
-        # Próbáljuk meg device paraméterrel létrehozni (újabb verziók ezt elfogadják)
+        # A linket változóba tesszük, hogy ne csússzon szét a sor másolásnál
+        gfpgan_url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth'
+
         try:
             face_enhancer = GFPGANer(
-                model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+                model_path=gfpgan_url,
                 upscale=args.outscale,
                 arch='clean',
                 channel_multiplier=2,
                 bg_upsampler=upsampler,
-                device=device  # ha a GFPGANer támogatja
+                device=device
             )
         except TypeError:
-            # ha a konstruktor nem fogadja a device parametert, fallback a régebbi inicializációra
             face_enhancer = GFPGANer(
-                model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+                model_path=gfpgan_url,
                 upscale=args.outscale,
                 arch='clean',
                 channel_multiplier=2,
                 bg_upsampler=upsampler
             )
 
-        # most KÉZZEL mozgassuk GPU-ra az összes belső modult
-        move_obj_to_device(face_enhancer, device, use_half)
-
-        # extra ellenőrzés: írjuk ki az első paraméter device-át, ha van
-        try:
-            # megtaláljuk az első modulparamétert
-            found = False
-            import torch.nn as nn
-            for attr in vars(face_enhancer).values():
-                if isinstance(attr, nn.Module):
-                    for p in attr.parameters(recurse=True):
-                        print("GFPGAN első paraméter device:", p.device)
-                        found = True
-                        break
-                if found:
-                    break
-        except Exception:
-            pass
+        # !!! JAVÍTÁS !!!
+        # Kényszerítjük a GFPGAN-t, hogy maradjon FP32-ben (False)
+        move_obj_to_device(face_enhancer, device, use_half=False)
 
     os.makedirs(args.output, exist_ok=True)
 
-    # beolvasási lista
     if os.path.isfile(args.input):
         paths = [args.input]
     else:
@@ -256,14 +232,13 @@ def main():
 
         try:
             if args.face_enhance and face_enhancer is not None:
-                # GFPGANer.enhance általában (np array) bemenetet vár
+                # GFPGANer futtatása
                 _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
             else:
                 output, _ = upsampler.enhance(img, outscale=args.outscale)
         except RuntimeError as error:
             print('Error during enhancement:', error)
-            print('If you encounter CUDA out of memory, try --tile with a smaller number or reduce --outscale.')
-            # tisztítás GPU memóriából
+            print('Trying to recover GPU memory...')
             try:
                 torch.cuda.empty_cache()
                 gc.collect()
@@ -284,7 +259,6 @@ def main():
             cv2.imwrite(save_path, output)
             print('Saved to', save_path)
 
-    # végső takarítás
     try:
         torch.cuda.empty_cache()
         gc.collect()
@@ -293,4 +267,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()