Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import gradio as gr | |
| import AnimeGANv3_src | |
| import numpy as np | |
| import logging | |
| import onnxruntime as ort | |
| import tempfile | |
| import threading | |
| # Logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s - %(levelname)s - %(message)s" | |
| ) | |
| # Check ONNX Runtime providers | |
| AVAILABLE_PROVIDERS = ort.get_available_providers() | |
| logging.info( | |
| "Available ONNX Runtime providers: %s", | |
| AVAILABLE_PROVIDERS | |
| ) | |
| if "CUDAExecutionProvider" in AVAILABLE_PROVIDERS: | |
| logging.info("CUDA provider is available.") | |
| else: | |
| logging.warning( | |
| "CUDA provider is NOT available. " | |
| "The model may run on CPU." | |
| ) | |
| class AnimeGANv3: | |
| def __init__(self): | |
| os.makedirs("output", exist_ok=True) | |
| os.makedirs("frames", exist_ok=True) | |
| self.lock = threading.Lock() | |
| logging.info("AnimeGANv3 initialized.") | |
| def process_frame(self, frame, style_code, det_face): | |
| """ | |
| Convert a single OpenCV BGR frame | |
| using the existing AnimeGANv3 module. | |
| """ | |
| frame_rgb = cv2.cvtColor( | |
| frame, | |
| cv2.COLOR_BGR2RGB | |
| ) | |
| output = AnimeGANv3_src.Convert( | |
| frame_rgb, | |
| style_code, | |
| det_face | |
| ) | |
| if output is None: | |
| raise RuntimeError( | |
| "AnimeGANv3 returned an empty frame." | |
| ) | |
| if not isinstance(output, np.ndarray): | |
| output = np.asarray(output) | |
| if output.ndim != 3 or output.shape[2] != 3: | |
| raise RuntimeError( | |
| f"Unexpected output shape: {output.shape}" | |
| ) | |
| # RGB -> BGR | |
| return cv2.cvtColor( | |
| output, | |
| cv2.COLOR_RGB2BGR | |
| ) | |
| def inference(self, video_path, style, if_face="No"): | |
| if not video_path: | |
| raise gr.Error("Please upload a video first.") | |
| style_codes = { | |
| "AnimeGANv3_Arcane": "A", | |
| "AnimeGANv3_Trump v1.0": "T", | |
| "AnimeGANv3_Shinkai": "S", | |
| "AnimeGANv3_PortraitSketch": "P", | |
| "AnimeGANv3_Hayao": "H", | |
| "AnimeGANv3_Disney v1.0": "D", | |
| "AnimeGANv3_JP_face v1.0": "J", | |
| "AnimeGANv3_Kpop v2.0": "K", | |
| } | |
| style_code = style_codes.get(style) | |
| if style_code is None: | |
| raise gr.Error("Invalid AnimeGANv3 style.") | |
| det_face = if_face == "Yes" | |
| logging.info( | |
| "Starting video conversion: style=%s, face_detection=%s", | |
| style, | |
| det_face | |
| ) | |
| with self.lock: | |
| cap = None | |
| out = None | |
| save_path = None | |
| try: | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| raise RuntimeError( | |
| "Could not open the uploaded video." | |
| ) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| if not np.isfinite(fps) or fps <= 0: | |
| fps = 25.0 | |
| frame_count = int( | |
| cap.get(cv2.CAP_PROP_FRAME_COUNT) | |
| ) | |
| width = int( | |
| cap.get(cv2.CAP_PROP_FRAME_WIDTH) | |
| ) | |
| height = int( | |
| cap.get(cv2.CAP_PROP_FRAME_HEIGHT) | |
| ) | |
| if width <= 0 or height <= 0: | |
| raise RuntimeError( | |
| "Could not determine video dimensions." | |
| ) | |
| logging.info( | |
| "Input: %s x %s, %.2f FPS, %s frames", | |
| width, | |
| height, | |
| fps, | |
| frame_count | |
| ) | |
| # Create a unique output filename. | |
| fd, save_path = tempfile.mkstemp( | |
| suffix=".mp4", | |
| prefix="animeganv3_", | |
| dir="output" | |
| ) | |
| os.close(fd) | |
| fourcc = cv2.VideoWriter_fourcc( | |
| *"mp4v" | |
| ) | |
| out = cv2.VideoWriter( | |
| save_path, | |
| fourcc, | |
| fps, | |
| (width, height) | |
| ) | |
| if not out.isOpened(): | |
| raise RuntimeError( | |
| "Could not initialize video writer." | |
| ) | |
| processed = 0 | |
| # Read and process frames immediately. | |
| # Avoid storing batches of frames in RAM. | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| stylized_frame = self.process_frame( | |
| frame, | |
| style_code, | |
| det_face | |
| ) | |
| # Preserve original video dimensions. | |
| if ( | |
| stylized_frame.shape[1] != width | |
| or stylized_frame.shape[0] != height | |
| ): | |
| stylized_frame = cv2.resize( | |
| stylized_frame, | |
| (width, height), | |
| interpolation=cv2.INTER_LINEAR | |
| ) | |
| out.write(stylized_frame) | |
| processed += 1 | |
| # Log progress occasionally, not every frame. | |
| if processed % 100 == 0: | |
| logging.info( | |
| "Processed %s / %s frames", | |
| processed, | |
| frame_count | |
| ) | |
| if processed == 0: | |
| raise RuntimeError( | |
| "No video frames were processed." | |
| ) | |
| logging.info( | |
| "Conversion complete: %s frames", | |
| processed | |
| ) | |
| return save_path | |
| except gr.Error: | |
| raise | |
| except Exception as error: | |
| logging.exception( | |
| "AnimeGANv3 conversion failed." | |
| ) | |
| if save_path and os.path.exists(save_path): | |
| try: | |
| os.remove(save_path) | |
| except OSError: | |
| pass | |
| raise gr.Error( | |
| f"Video conversion failed: {error}" | |
| ) | |
| finally: | |
| if cap is not None: | |
| cap.release() | |
| if out is not None: | |
| out.release() | |
| anime_gan = AnimeGANv3() | |
| title = "AnimeGANv3: Video to Anime Converter" | |
| description = r""" | |
| Upload a video to convert it into anime style using AnimeGANv3.<br> | |
| Select a style and choose whether to optimize for faces.<br> | |
| Credits to fine tuner Asher_Chan<br> | |
| <a href="https://github.com/TachibanaYoshino/AnimeGANv3" | |
| target="_blank"><b>AnimeGANv3 GitHub</b></a> | | |
| <a href="https://www.patreon.com/Asher_Chan" | |
| target="_blank"><b>Patreon</b></a> | |
| """ | |
| iface = gr.Interface( | |
| fn=anime_gan.inference, | |
| inputs=[ | |
| gr.Video(label="Input Video"), | |
| gr.Dropdown( | |
| choices=[ | |
| "AnimeGANv3_Hayao", | |
| "AnimeGANv3_Shinkai", | |
| "AnimeGANv3_Arcane", | |
| "AnimeGANv3_Trump v1.0", | |
| "AnimeGANv3_Disney v1.0", | |
| "AnimeGANv3_PortraitSketch", | |
| "AnimeGANv3_JP_face v1.0", | |
| "AnimeGANv3_Kpop v2.0", | |
| ], | |
| label="AnimeGANv3 Style", | |
| value="AnimeGANv3_Arcane" | |
| ), | |
| gr.Radio( | |
| choices=["Yes", "No"], | |
| label="Extract face", | |
| value="No" | |
| ), | |
| ], | |
| outputs=[ | |
| gr.Video(label="Output Video") | |
| ], | |
| title=title, | |
| description=description, | |
| allow_flagging="never" | |
| ) | |
| if __name__ == "__main__": | |
| iface.queue().launch() |