Update inference.py

inference.py

@@ -1,362 +1,383 @@
-from os import listdir, path
+import os
+import argparse
+import subprocess
+import platform
+
 import numpy as np
-import scipy, cv2, os, sys, argparse, audio
-import json, subprocess, random, string
+import cv2
+import torch
 from tqdm import tqdm
-from glob import glob
-import torch, face_detection
+
+from face_detection import FaceAlignment, LandmarksType
 from wav2lip_models import Wav2Lip
-import platform
 from face_parsing import init_parser, swap_regions
-from esrgan.upsample import upscale
-from esrgan.upsample import load_sr
-from basicsr.archs.rrdbnet_arch import RRDBNet
+from esrgan.upsample import upscale, load_sr
 from basicsr.utils.download_util import load_file_from_url
 
-parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
-
-parser.add_argument('--checkpoint_path', type=str, default="checkpoints/wav2lip_gan.pth",
-					help='Name of saved checkpoint to load weights from', required=False)
-
-parser.add_argument('--segmentation_path', type=str, default="checkpoints/face_segmentation.pth",
-					help='Name of saved checkpoint of segmentation network', required=False)
-
-parser.add_argument('--sr_path', type=str, default='weights/4x_BigFace_v3_Clear.pth', 
-					help='Name of saved checkpoint of super-resolution network', required=False)
-
-parser.add_argument('--face', type=str, 
-					help='Filepath of video/image that contains faces to use', required=True)
-parser.add_argument('--audio', type=str, 
-					help='Filepath of video/audio file to use as raw audio source', required=True)
-parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', 
-								default='results/result_voice.mp4')
-
-
-parser.add_argument('--static', type=bool, 
-					help='If True, then use only first video frame for inference', default=False)
-parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)', 
-					default=25., required=False)
-
-parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], 
-					help='Padding (top, bottom, left, right). Please adjust to include chin at least')
-
-parser.add_argument('--face_det_batch_size', type=int, 
-					help='Batch size for face detection', default=16)
-parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
-
-parser.add_argument('--resize_factor', default=1, type=int, 
-			help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
-
-parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], 
-					help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' 
-					'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
-
-parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], 
-					help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
-					'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
-
-parser.add_argument('--rotate', default=False, action='store_true',
-					help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
-					'Use if you get a flipped result, despite feeding a normal looking video')
-
-parser.add_argument('--nosmooth', default=False, action='store_true',
-					help='Prevent smoothing face detections over a short temporal window')
-parser.add_argument('--no_seg', default=False, action='store_true',
-					help='Prevent using face segmentation')
-parser.add_argument('--no_sr', default=False, action='store_true',
-					help='Prevent using super resolution')
-parser.add_argument('--enhance_face', default=None, choices=['gfpgan','codeformer'],
-					help='Use GFP-GAN to enhance facial details.')
-parser.add_argument('-w', '--fidelity_weight', type=float, default=0.75, 
-            				help='Balance the quality and fidelity. Default: 0.75')
-parser.add_argument('--save_frames', default=False, action='store_true',
-					help='Save each frame as an image. Use with caution')
-parser.add_argument('--gt_path', type=str, 
-					help='Where to store saved ground truth frames', required=False)
-parser.add_argument('--pred_path', type=str, 
-					help='Where to store frames produced by algorithm', required=False)
-parser.add_argument('--save_as_video', action="store_true", default=False,
-					help='Whether to save frames as video', required=False)
-parser.add_argument('--image_prefix', type=str, default="",
-					help='Prefix to save frames with', required=False)
-
-args = parser.parse_args()
-args.img_size = 96
-
-if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
-	args.static = True
+import audio  # Предполагается, что это ваш модуль для обработки аудио
+
+# Оптимизированные импорты: удалены неиспользуемые библиотеки (scipy, json, random, string, glob)
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
+    
+    parser.add_argument('--checkpoint_path', type=str, default="checkpoints/wav2lip_gan.pth",
+                        help='Name of saved checkpoint to load weights from', required=False)
+    
+    parser.add_argument('--segmentation_path', type=str, default="checkpoints/face_segmentation.pth",
+                        help='Name of saved checkpoint of segmentation network', required=False)
+    
+    parser.add_argument('--sr_path', type=str, default='weights/4x_BigFace_v3_Clear.pth', 
+                        help='Name of saved checkpoint of super-resolution network', required=False)
+    
+    parser.add_argument('--face', type=str, 
+                        help='Filepath of video/image that contains faces to use', required=True)
+    parser.add_argument('--audio', type=str, 
+                        help='Filepath of video/audio file to use as raw audio source', required=True)
+    parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', 
+                                    default='results/result_voice.mp4')
+    
+    parser.add_argument('--static', action='store_true', 
+                        help='If set, use only first video frame for inference')
+    parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)', 
+                        default=25., required=False)
+    
+    parser.add_argument('--pads', nargs=4, type=int, default=[0, 10, 0, 0], 
+                        help='Padding (top, bottom, left, right). Please adjust to include chin at least')
+    
+    parser.add_argument('--face_det_batch_size', type=int, 
+                        help='Batch size for face detection', default=16)
+    parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
+    
+    parser.add_argument('--resize_factor', default=1, type=int, 
+                        help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
+    
+    parser.add_argument('--crop', nargs=4, type=int, default=[0, -1, 0, -1], 
+                        help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' 
+                             'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
+    
+    parser.add_argument('--box', nargs=4, type=int, default=[-1, -1, -1, -1], 
+                        help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
+                             'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
+    
+    parser.add_argument('--rotate', action='store_true',
+                        help='Sometimes videos taken from a phone can be flipped 90deg. If set, will flip video right by 90deg.'
+                             'Use if you get a flipped result, despite feeding a normal looking video')
+    
+    parser.add_argument('--nosmooth', action='store_true',
+                        help='Prevent smoothing face detections over a short temporal window')
+    parser.add_argument('--no_seg', action='store_true',
+                        help='Prevent using face segmentation')
+    parser.add_argument('--no_sr', action='store_true',
+                        help='Prevent using super resolution')
+    parser.add_argument('--enhance_face', choices=['gfpgan','codeformer'],
+                        help='Use GFP-GAN or CodeFormer to enhance facial details.')
+    parser.add_argument('-w', '--fidelity_weight', type=float, default=0.75, 
+                        help='Balance the quality and fidelity. Default: 0.75')
+    parser.add_argument('--save_frames', action='store_true',
+                        help='Save each frame as an image. Use with caution')
+    parser.add_argument('--gt_path', type=str, 
+                        help='Where to store saved ground truth frames', required=False)
+    parser.add_argument('--pred_path', type=str, 
+                        help='Where to store frames produced by algorithm', required=False)
+    parser.add_argument('--save_as_video', action="store_true", default=False,
+                        help='Whether to save frames as video', required=False)
+    parser.add_argument('--image_prefix', type=str, default="",
+                        help='Prefix to save frames with', required=False)
+    
+    args = parser.parse_args()
+
+    # Определение, является ли ввод статичным изображением
+    if os.path.isfile(args.face) and os.path.splitext(args.face)[1].lower() in ['.jpg', '.png', '.jpeg']:
+        args.static = True
+
+    args.img_size = 96
+    return args
+
 
 def get_smoothened_boxes(boxes, T):
-	for i in range(len(boxes)):
-		if i + T > len(boxes):
-			window = boxes[len(boxes) - T:]
-		else:
-			window = boxes[i : i + T]
-		boxes[i] = np.mean(window, axis=0)
-	return boxes
-
-def face_detect(images):
-	detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, 
-											flip_input=False, device=device)
-
-	batch_size = args.face_det_batch_size
-	
-	while 1:
-		predictions = []
-		try:
-			for i in range(0, len(images), batch_size):
-				predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
-		except RuntimeError:
-			if batch_size == 1: 
-				raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
-			batch_size //= 2
-			print('Recovering from OOM error; New batch size: {}'.format(batch_size))
-			continue
-		break
-
-	results = []
-	pady1, pady2, padx1, padx2 = args.pads
-	for rect, image in zip(predictions, images):
-		if rect is None:
-			continue
-		y1 = max(0, rect[1] - pady1)
-		y2 = min(image.shape[0], rect[3] + pady2)
-		x1 = max(0, rect[0] - padx1)
-		x2 = min(image.shape[1], rect[2] + padx2)
-		
-		results.append([x1, y1, x2, y2])
-
-	boxes = np.array(results)
-	if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
-	results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
-
-	del detector
-	return results 
-
-def datagen(mels):
+    for i in range(len(boxes)):
+        window = boxes[max(i - T + 1, 0):i + 1]
+        boxes[i] = np.mean(window, axis=0)
+    return boxes
+
+
+def face_detect(detector, images, args):
+    predictions = []
+    batch_size = args.face_det_batch_size
+    
+    try:
+        for i in range(0, len(images), batch_size):
+            batch_images = np.array(images[i:i + batch_size])
+            predictions.extend(detector.get_detections_for_batch(batch_images))
+    except RuntimeError:
+        if batch_size == 1: 
+            raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
+        batch_size //= 2
+        print(f'Recovering from OOM error; New batch size: {batch_size}')
+        return face_detect(detector, images, args)  # Рекурсивный вызов с уменьшенным batch_size
+
+    results = []
+    pady1, pady2, padx1, padx2 = args.pads
+    for rect, image in zip(predictions, images):
+        if rect is None:
+            continue
+        y1 = max(0, rect[1] - pady1)
+        y2 = min(image.shape[0], rect[3] + pady2)
+        x1 = max(0, rect[0] - padx1)
+        x2 = min(image.shape[1], rect[2] + padx2)
+        
+        results.append([x1, y1, x2, y2])
+
+    boxes = np.array(results)
+    if not args.nosmooth and len(boxes) > 0:
+        boxes = get_smoothened_boxes(boxes, T=5)
+    
+    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] 
+               for image, (x1, y1, x2, y2) in zip(images, boxes)]
+    
+    return results 
+
+
+def datagen(mels, reader, detector, args):
     img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
 
-    # Uncommented code removed for clarity
-
-    reader = read_frames()
-
-    for i, m in enumerate(mels):
-        try:
-            frame_to_save = next(reader)
-        except StopIteration:
-            reader = read_frames()
+    for m in mels:
+        frame_to_save = next(reader, None)
+        if frame_to_save is None:
+            reader = read_frames(args.face, args.resize_factor, args.rotate, args.crop)
             frame_to_save = next(reader, None)
-
-        if frame_to_save is not None:
-            face_detect_result = face_detect([frame_to_save])
-            if len(face_detect_result) > 0:  # Check if face detection was successful
-                face, coords = face_detect_result[0]
-                face = cv2.resize(face, (args.img_size, args.img_size))
-                img_batch.append(face)
-                mel_batch.append(m)
-                frame_batch.append(frame_to_save)
-                coords_batch.append(coords)
+            if frame_to_save is None:
+                break
+
+        face_detect_result = face_detect(detector, [frame_to_save], args)
+        if len(face_detect_result) > 0:  # Check if face detection was successful
+            face, coords = face_detect_result[0]
+            face = cv2.resize(face, (args.img_size, args.img_size))
+            img_batch.append(face)
+            mel_batch.append(m)
+            frame_batch.append(frame_to_save)
+            coords_batch.append(coords)
 
         if len(img_batch) >= args.wav2lip_batch_size:
-            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+            img_batch_np = np.asarray(img_batch)
+            mel_batch_np = np.asarray(mel_batch)
 
-            img_masked = img_batch.copy()
+            img_masked = img_batch_np.copy()
             img_masked[:, args.img_size // 2:] = 0
 
-            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+            img_batch_np = np.concatenate((img_masked, img_batch_np), axis=3) / 255.0
+            mel_batch_np = mel_batch_np.reshape(len(mel_batch_np), mel_batch_np.shape[1], mel_batch_np.shape[2], 1)
 
-            yield img_batch, mel_batch, frame_batch, coords_batch
+            yield img_batch_np, mel_batch_np, frame_batch, coords_batch
             img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
 
     if len(img_batch) > 0:
-        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+        img_batch_np = np.asarray(img_batch)
+        mel_batch_np = np.asarray(mel_batch)
 
-        img_masked = img_batch.copy()
+        img_masked = img_batch_np.copy()
         img_masked[:, args.img_size // 2:] = 0
 
-        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
-
-        yield img_batch, mel_batch, frame_batch, coords_batch
-
-mel_step_size = 16
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print('Using {} for inference.'.format(device))
-
-def _load(checkpoint_path):
-	if device == 'cuda':
-		checkpoint = torch.load(checkpoint_path)
-	else:
-		checkpoint = torch.load(checkpoint_path,
-								map_location=lambda storage, loc: storage)
-	return checkpoint
-
-def load_model(path):
-	model = Wav2Lip()
-	print("Load checkpoint from: {}".format(path))
-	checkpoint = _load(path)
-	s = checkpoint["state_dict"]
-	new_s = {}
-	for k, v in s.items():
-		new_s[k.replace('module.', '')] = v
-	model.load_state_dict(new_s)
-
-	model = model.to(device)
-	return model.eval()
-
-def read_frames():
-	if args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
-		face = cv2.imread(args.face)
-		while 1:
-			yield face
-
-	video_stream = cv2.VideoCapture(args.face)
-	fps = video_stream.get(cv2.CAP_PROP_FPS)
-
-	print('Reading video frames from start...')
-
-	while 1:
-		still_reading, frame = video_stream.read()
-		if not still_reading:
-			video_stream.release()
-			break
-		if args.resize_factor > 1:
-			frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
-
-		if args.rotate:
-			frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
-
-		y1, y2, x1, x2 = args.crop
-		if x2 == -1: x2 = frame.shape[1]
-		if y2 == -1: y2 = frame.shape[0]
-
-		frame = frame[y1:y2, x1:x2]
+        img_batch_np = np.concatenate((img_masked, img_batch_np), axis=3) / 255.0
+        mel_batch_np = mel_batch_np.reshape(len(mel_batch_np), mel_batch_np.shape[1], mel_batch_np.shape[2], 1)
+
+        yield img_batch_np, mel_batch_np, frame_batch, coords_batch
+
+
+def load_checkpoint(checkpoint_path, device):
+    if device == 'cuda':
+        checkpoint = torch.load(checkpoint_path)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+    return checkpoint
+
+
+def load_model(checkpoint_path, device):
+    model = Wav2Lip()
+    print(f"Loading checkpoint from: {checkpoint_path}")
+    checkpoint = load_checkpoint(checkpoint_path, device)
+    state_dict = checkpoint["state_dict"]
+    new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+    model.load_state_dict(new_state_dict)
+    model = model.to(device)
+    model.eval()
+    return model
+
+
+def read_frames(face_path, resize_factor, rotate, crop):
+    if os.path.splitext(face_path)[1].lower() in ['.jpg', '.png', '.jpeg']:
+        face = cv2.imread(face_path)
+        if resize_factor > 1:
+            face = cv2.resize(face, (face.shape[1]//resize_factor, face.shape[0]//resize_factor))
+        if rotate:
+            face = cv2.rotate(face, cv2.ROTATE_90_CLOCKWISE)
+        y1, y2, x1, x2 = crop
+        if x2 == -1: x2 = face.shape[1]
+        if y2 == -1: y2 = face.shape[0]
+        face = face[y1:y2, x1:x2]
+        while True:
+            yield face
+    else:
+        video_stream = cv2.VideoCapture(face_path)
+        fps = video_stream.get(cv2.CAP_PROP_FPS)
+        print('Reading video frames from start...')
+
+        while True:
+            still_reading, frame = video_stream.read()
+            if not still_reading:
+                video_stream.release()
+                break
+            if resize_factor > 1:
+                frame = cv2.resize(frame, (frame.shape[1]//resize_factor, frame.shape[0]//resize_factor))
+            if rotate:
+                frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
+            y1, y2, x1, x2 = crop
+            if x2 == -1: x2 = frame.shape[1]
+            if y2 == -1: y2 = frame.shape[0]
+            frame = frame[y1:y2, x1:x2]
+            yield frame
 
-		yield frame
 
 def main():
-	if not os.path.isfile(args.face):
-		raise ValueError('--face argument must be a valid path to video/image file')
-
-	elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
-		fps = args.fps
-	else:
-		video_stream = cv2.VideoCapture(args.face)
-		fps = video_stream.get(cv2.CAP_PROP_FPS)
-		video_stream.release()
-
-
-	if not args.audio.endswith('.wav'):
-		print('Extracting raw audio...')
-		command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav')
-
-		subprocess.call(command, shell=True)
-		args.audio = 'temp/temp.wav'
-
-	wav = audio.load_wav(args.audio, 16000)
-	mel = audio.melspectrogram(wav)
-	print(mel.shape)
-
-	if np.isnan(mel.reshape(-1)).sum() > 0:
-		raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
-
-	mel_chunks = []
-	mel_idx_multiplier = 80./fps 
-	i = 0
-	while 1:
-		start_idx = int(i * mel_idx_multiplier)
-		if start_idx + mel_step_size > len(mel[0]):
-			mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
-			break
-		mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
-		i += 1
-
-	print("Length of mel chunks: {}".format(len(mel_chunks)))
-
-	batch_size = args.wav2lip_batch_size
-	gen = datagen(mel_chunks)
-
-
-
-	if args.save_as_video:
-		gt_out = cv2.VideoWriter("temp/gt.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (384, 384))
-		pred_out = cv2.VideoWriter("temp/pred.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (96, 96))
-
-	abs_idx = 0
-	for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, 
-											total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
-		if i == 0:
-			if not args.no_seg==True:
-				print("Loading segmentation network...")
-				seg_net = load_file_from_url(
-					url='https://github.com/GucciFlipFlops1917/wav2lip-hq-updated-ESRGAN/releases/download/v0.0.1/face_segmentation.pth',
-					model_dir='checkpoints', progress=True, file_name=None)
-				seg_net = init_parser(args.segmentation_path)
-			if not args.no_sr==True:
-				print("Loading super resolution model...")
-				run_params = load_sr(args.sr_path, device, args.enhance_face)
-			
-			model_path = load_file_from_url(
-				url='https://github.com/GucciFlipFlops1917/wav2lip-hq-updated-ESRGAN/releases/download/v0.0.1/wav2lip_gan.pth',
-				model_dir='checkpoints', progress=True, file_name=None)
-			model = load_model(args.checkpoint_path)
-			print ("Model loaded")
-
-			frame_h, frame_w = next(read_frames()).shape[:-1]
-			out = cv2.VideoWriter('temp/result.avi', 
-									cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
-
-		img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
-		mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
-
-		with torch.no_grad():
-			pred = model(mel_batch, img_batch)
-
-		pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
-		
-		for p, f, c in zip(pred, frames, coords):
-			y1, y2, x1, x2 = c
-
-			if args.save_frames:
-				if args.save_as_video:
-					pred_out.write(p.astype(np.uint8))
-					gt_out.write(cv2.resize(f[y1:y2, x1:x2], (384, 384)))
-				else:
-					cv2.imwrite(f"{args.gt_path}/{args.image_prefix}{abs_idx}.png", f[y1:y2, x1:x2])
-					cv2.imwrite(f"{args.pred_path}/{args.image_prefix}{abs_idx}.png", p)
-					abs_idx += 1
-
-			if not args.no_sr:
-				if args.enhance_face==None:
-					p = upscale(p, 0, run_params)
-				elif args.enhance_face=='codeformer':
-					p = upscale(p, 2, [run_params, device, args.fidelity_weight])
-				elif args.enhance_face=='gfpgan':
-					p = upscale(p, 1, run_params)
-			p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
-			
-			if args.no_seg==False:
-				p = swap_regions(f[y1:y2, x1:x2], p, seg_net)
-
-			f[y1:y2, x1:x2] = p
-			out.write(f)
-
-	out.release()
-
-	command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile)
-	subprocess.call(command, shell=platform.system() != 'Windows')
-
-	if args.save_frames and args.save_as_video:
-		gt_out.release()
-		pred_out.release()
-
-		command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/gt.avi', args.gt_path)
-		subprocess.call(command, shell=platform.system() != 'Windows')
-
-		command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/pred.avi', args.pred_path)
-		subprocess.call(command, shell=platform.system() != 'Windows')
+    args = parse_arguments()
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f'Using {device} for inference.')
+
+    # Инициализация моделей вне циклов
+    detector = FaceAlignment(LandmarksType._2D, flip_input=False, device=device)
+
+    if not args.no_seg:
+        print("Loading segmentation network...")
+        seg_net = init_parser(args.segmentation_path)
+    else:
+        seg_net = None
+
+    if not args.no_sr:
+        print("Loading super resolution model...")
+        run_params = load_sr(args.sr_path, device, args.enhance_face)
+    else:
+        run_params = None
+
+    model = load_model(args.checkpoint_path, device)
+    print("Model loaded")
+
+    if not os.path.isfile(args.face):
+        raise ValueError('--face argument must be a valid path to video/image file')
+
+    if not args.audio.endswith('.wav'):
+        print('Extracting raw audio...')
+        temp_wav = 'temp/temp.wav'
+        command = f'ffmpeg -y -i "{args.audio}" -strict -2 "{temp_wav}"'
+        subprocess.call(command, shell=True)
+        args.audio = temp_wav
+
+    wav = audio.load_wav(args.audio, 16000)
+    mel = audio.melspectrogram(wav)
+    print(mel.shape)
+
+    if np.isnan(mel).any():
+        raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
+
+    mel_step_size = 16
+    fps = args.fps if args.static else None
+
+    if not args.static:
+        video_stream = cv2.VideoCapture(args.face)
+        fps = video_stream.get(cv2.CAP_PROP_FPS)
+        video_stream.release()
+
+    mel_idx_multiplier = 80.0 / fps 
+    mel_chunks = []
+    i = 0
+    while True:
+        start_idx = int(i * mel_idx_multiplier)
+        if start_idx + mel_step_size > mel.shape[1]:
+            mel_chunks.append(mel[:, -mel_step_size:])
+            break
+        mel_chunks.append(mel[:, start_idx:start_idx + mel_step_size])
+        i += 1
+
+    print(f"Length of mel chunks: {len(mel_chunks)}")
+
+    reader = read_frames(args.face, args.resize_factor, args.rotate, args.crop)
+    generator = datagen(mel_chunks, reader, detector, args)
+
+    if args.save_as_video:
+        frame_sample = next(reader)
+        frame_h, frame_w = frame_sample.shape[:2]
+        out = cv2.VideoWriter('temp/result.avi', 
+                              cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
+        if args.save_frames:
+            gt_out = cv2.VideoWriter("temp/gt.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (384, 384))
+            pred_out = cv2.VideoWriter("temp/pred.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (96, 96))
+    else:
+        out = None
+        gt_out = None
+        pred_out = None
+
+    abs_idx = 0
+    for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(generator, 
+                                        total=int(np.ceil(len(mel_chunks)/args.wav2lip_batch_size)))):
+        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
+        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
+
+        with torch.no_grad():
+            pred = model(mel_batch, img_batch)
+
+        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.0
+        
+        for p, f, c in zip(pred, frames, coords):
+            y1, y2, x1, x2 = c
+
+            if args.save_frames:
+                if args.save_as_video:
+                    pred_out.write(p.astype(np.uint8))
+                    gt_resized = cv2.resize(f[y1:y2, x1:x2], (384, 384))
+                    gt_out.write(gt_resized)
+                else:
+                    if args.gt_path and args.pred_path:
+                        os.makedirs(args.gt_path, exist_ok=True)
+                        os.makedirs(args.pred_path, exist_ok=True)
+                        cv2.imwrite(f"{args.gt_path}/{args.image_prefix}{abs_idx}.png", f[y1:y2, x1:x2])
+                        cv2.imwrite(f"{args.pred_path}/{args.image_prefix}{abs_idx}.png", p)
+                        abs_idx += 1
+
+            if not args.no_sr:
+                if args.enhance_face is None:
+                    p = upscale(p, 0, run_params)
+                elif args.enhance_face == 'codeformer':
+                    p = upscale(p, 2, [run_params, device, args.fidelity_weight])
+                elif args.enhance_face == 'gfpgan':
+                    p = upscale(p, 1, run_params)
+
+            p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
+            
+            if not args.no_seg and seg_net is not None:
+                p = swap_regions(f[y1:y2, x1:x2], p, seg_net)
+
+            f[y1:y2, x1:x2] = p
+            if out:
+                out.write(f)
+
+    if out:
+        out.release()
+
+    # Объединение аудио и видео
+    final_command = f'ffmpeg -y -i "{args.audio}" -i "temp/result.avi" -strict -2 -q:v 1 "{args.outfile}"'
+    subprocess.call(final_command, shell=(platform.system() != 'Windows'))
+
+    if args.save_frames and args.save_as_video:
+        gt_out.release()
+        pred_out.release()
+
+        gt_video_cmd = f'ffmpeg -y -i "temp/gt.avi" -i "{args.audio}" -strict -2 -q:v 1 "{args.gt_path}"'
+        pred_video_cmd = f'ffmpeg -y -i "temp/pred.avi" -i "{args.audio}" -strict -2 -q:v 1 "{args.pred_path}"'
+
+        subprocess.call(gt_video_cmd, shell=(platform.system() != 'Windows'))
+        subprocess.call(pred_video_cmd, shell=(platform.system() != 'Windows'))
 
 
 if __name__ == '__main__':
-	main()
+    main()