Remove broken .so file from repo and LFS

34bd296 over 1 year ago

3.09 kB

import os
import torch
import shutil
import uuid
from pathlib import Path
import warnings
from typing import BinaryIO
from src.wav2lip.models import Wav2Lip
from src.wav2lip import inference

warnings.filterwarnings("ignore")

# Directories setup
UPLOAD_DIR = Path("uploads")
OUTPUT_DIR = Path("outputs")
RESULT_TEMP = Path("wav2lip/results/result_voice.mp4")
CHECKPOINT_PATH = "wav2lip_checkpoints/wav2lip_gan.pth"

UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
Path("wav2lip/results").mkdir(parents=True, exist_ok=True)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = None


def load_model(checkpoint_path: str):
    global MODEL
    if MODEL is not None:
        return MODEL

    print(f"Loading Wav2Lip model from {checkpoint_path} on {DEVICE}")
    model = Wav2Lip()

    # Load checkpoint on the selected device
    checkpoint = torch.load(checkpoint_path, map_location=DEVICE)

    # Fix possible 'module.' prefix in keys if loaded from multi-GPU training
    state_dict = {
        k.replace("module.", ""): v for k, v in checkpoint["state_dict"].items()
    }
    model.load_state_dict(state_dict)

    # Move model to device (GPU or CPU)
    model = model.to(DEVICE)
    model.eval()

    MODEL = model
    return MODEL


def unload_model():
    global MODEL
    if MODEL:
        del MODEL
        MODEL = None
        if DEVICE == "mps":
            torch.mps.empty_cache()
        print("Model unloaded and memory freed.")


def predict(image_file: BinaryIO, audio_file: BinaryIO, model) -> str:
    """
    Receives image/audio files as file-like objects (e.g. BytesIO), saves them, runs prediction,
    and returns the path to the result video.
    """
    # Generate session-specific filenames
    session_id = str(uuid.uuid4())
    image_path = UPLOAD_DIR / f"{session_id}.jpg"
    audio_path = UPLOAD_DIR / f"{session_id}.mp3"
    output_path = OUTPUT_DIR / f"{session_id}_result.mp4"

    # Save input files
    with open(image_path, "wb") as f:
        shutil.copyfileobj(image_file, f)
    with open(audio_path, "wb") as f:
        shutil.copyfileobj(audio_file, f)

    print(f"Saved inputs to: {image_path}, {audio_path}")

    # # Load the model
    # model = load_model(CHECKPOINT_PATH)

    # Run inference
    inference.main(str(image_path), str(audio_path), model)

    # Move output
    if not RESULT_TEMP.exists():
        raise RuntimeError("Inference failed. Output video not found.")

    shutil.move(RESULT_TEMP, output_path)
    print(f"Saved output to: {output_path}")

    return str(output_path)


# from pathlib import Path
# from wav2lip_api import predict, load_model  # Replace with your actual script/module name

# # Load image and audio files as binary
# image_path = "sample_inputs/sample.jpg"  # Replace with actual image path
# audio_path = "sample_inputs/sample.mp3"  # Replace with actual audio path

# with open(image_path, "rb") as img_file, open(audio_path, "rb") as aud_file:
#     result_path = predict(img_file, aud_file)

# print(f"Result video saved at: {result_path}")

	import os
	import torch
	import shutil
	import uuid
	from pathlib import Path
	import warnings
	from typing import BinaryIO
	from src.wav2lip.models import Wav2Lip
	from src.wav2lip import inference

	warnings.filterwarnings("ignore")

	# Directories setup
	UPLOAD_DIR = Path("uploads")
	OUTPUT_DIR = Path("outputs")
	RESULT_TEMP = Path("wav2lip/results/result_voice.mp4")
	CHECKPOINT_PATH = "wav2lip_checkpoints/wav2lip_gan.pth"

	UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
	OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
	Path("wav2lip/results").mkdir(parents=True, exist_ok=True)

	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
	MODEL = None


	def load_model(checkpoint_path: str):
	global MODEL
	if MODEL is not None:
	return MODEL

	print(f"Loading Wav2Lip model from {checkpoint_path} on {DEVICE}")
	model = Wav2Lip()

	# Load checkpoint on the selected device
	checkpoint = torch.load(checkpoint_path, map_location=DEVICE)

	# Fix possible 'module.' prefix in keys if loaded from multi-GPU training
	state_dict = {
	k.replace("module.", ""): v for k, v in checkpoint["state_dict"].items()
	}
	model.load_state_dict(state_dict)

	# Move model to device (GPU or CPU)
	model = model.to(DEVICE)
	model.eval()

	MODEL = model
	return MODEL


	def unload_model():
	global MODEL
	if MODEL:
	del MODEL
	MODEL = None
	if DEVICE == "mps":
	torch.mps.empty_cache()
	print("Model unloaded and memory freed.")


	def predict(image_file: BinaryIO, audio_file: BinaryIO, model) -> str:
	"""
	Receives image/audio files as file-like objects (e.g. BytesIO), saves them, runs prediction,
	and returns the path to the result video.
	"""
	# Generate session-specific filenames
	session_id = str(uuid.uuid4())
	image_path = UPLOAD_DIR / f"{session_id}.jpg"
	audio_path = UPLOAD_DIR / f"{session_id}.mp3"
	output_path = OUTPUT_DIR / f"{session_id}_result.mp4"

	# Save input files
	with open(image_path, "wb") as f:
	shutil.copyfileobj(image_file, f)
	with open(audio_path, "wb") as f:
	shutil.copyfileobj(audio_file, f)

	print(f"Saved inputs to: {image_path}, {audio_path}")

	# # Load the model
	# model = load_model(CHECKPOINT_PATH)

	# Run inference
	inference.main(str(image_path), str(audio_path), model)

	# Move output
	if not RESULT_TEMP.exists():
	raise RuntimeError("Inference failed. Output video not found.")

	shutil.move(RESULT_TEMP, output_path)
	print(f"Saved output to: {output_path}")

	return str(output_path)


	# from pathlib import Path
	# from wav2lip_api import predict, load_model # Replace with your actual script/module name

	# # Load image and audio files as binary
	# image_path = "sample_inputs/sample.jpg" # Replace with actual image path
	# audio_path = "sample_inputs/sample.mp3" # Replace with actual audio path

	# with open(image_path, "rb") as img_file, open(audio_path, "rb") as aud_file:
	# result_path = predict(img_file, aud_file)

	# print(f"Result video saved at: {result_path}")