Extras/ffff_assets/target.py

from typing import Optional, Tuple
import gradio
import os
import tempfile
import requests
from PIL import ImageGrab, Image

from ffff import state_manager, wording
from ffff.face_store import clear_reference_faces, clear_static_faces
from ffff.filesystem import get_file_size, is_image, is_video
from ffff.uis.core import register_ui_component
from ffff.vision import get_video_frame, normalize_frame_color
from ffff.uis.typing import File

FILE_SIZE_LIMIT = 512 * 1024 * 1024

TARGET_FILE: Optional[gradio.File] = None
TARGET_IMAGE: Optional[gradio.Image] = None
TARGET_VIDEO: Optional[gradio.Video] = None


def render() -> None:
    global TARGET_FILE
    global TARGET_IMAGE
    global TARGET_VIDEO

    # Obtener si el archivo es imagen o video
    is_target_image = is_image(state_manager.get_item('target_path'))
    is_target_video = is_video(state_manager.get_item('target_path'))

    TARGET_FILE = gradio.File(
        label=wording.get('uis.target_file'),
        file_count='single',
        file_types=['image', 'video'],
        value=state_manager.get_item('target_path') if is_target_image or is_target_video else None
    )

    target_image_options = {'show_label': False, 'visible': False}
    target_video_options = {'show_label': False, 'visible': False}

    if is_target_image:
        target_image_options['value'] = TARGET_FILE.value.get('path')
        target_image_options['visible'] = True

    if is_target_video:
        if get_file_size(state_manager.get_item('target_path')) > FILE_SIZE_LIMIT:
            preview_vision_frame = normalize_frame_color(get_video_frame(state_manager.get_item('target_path')))
            target_image_options['value'] = preview_vision_frame
            target_image_options['visible'] = True
        else:
            target_video_options['value'] = TARGET_FILE.value.get('path')
            target_video_options['visible'] = True

    TARGET_IMAGE = gradio.Image(**target_image_options)
    TARGET_VIDEO = gradio.Video(**target_video_options)

    register_ui_component('target_image', TARGET_IMAGE)
    register_ui_component('target_video', TARGET_VIDEO)


def listen() -> None:
    TARGET_FILE.change(fn=update, inputs=TARGET_FILE, outputs=[TARGET_IMAGE, TARGET_VIDEO])


def update(file: File) -> Tuple[gradio.Image, gradio.Video]:
    clear_reference_faces()
    clear_static_faces()
    if file and is_image(file.name):
        state_manager.set_item('target_path', file.name)
        return gradio.Image(value=file.name, visible=True), gradio.Video(value=None, visible=False)
    if file and is_video(file.name):
        state_manager.set_item('target_path', file.name)
        if get_file_size(file.name) > FILE_SIZE_LIMIT:
            preview_vision_frame = normalize_frame_color(get_video_frame(file.name))
            return gradio.Image(value=preview_vision_frame, visible=True), gradio.Video(value=None, visible=False)
        return gradio.Image(value=None, visible=False), gradio.Video(value=file.name, visible=True)
    state_manager.clear_item('target_path')
    return gradio.Image(value=None, visible=False), gradio.Video(value=None, visible=False)


# Función para descargar un archivo desde URL
def download_file_from_url(url: str) -> Optional[str]:
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(url)[-1])
        with open(temp_file.name, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
        return temp_file.name
    except Exception as e:
        print(f"Error al descargar el archivo: {e}")
        return None


# Función para manejar la entrada de URL en Target
def handle_url_input(url: str) -> Tuple[gradio.Image, gradio.Video]:
    downloaded_file_path = download_file_from_url(url)
    if downloaded_file_path:
        state_manager.set_item('target_path', downloaded_file_path)
        if is_image(downloaded_file_path):
            return gradio.Image(value=downloaded_file_path, visible=True), gradio.Video(value=None, visible=False)
        if is_video(downloaded_file_path):
            if get_file_size(downloaded_file_path) > FILE_SIZE_LIMIT:
                preview_vision_frame = normalize_frame_color(get_video_frame(downloaded_file_path))
                return gradio.Image(value=preview_vision_frame, visible=True), gradio.Video(value=None, visible=False)
            return gradio.Image(value=None, visible=False), gradio.Video(value=downloaded_file_path, visible=True)
    return gradio.Image(value=None, visible=False), gradio.Video(value=None, visible=False)


# Función para manejar pegado desde el portapapeles en Target
def handle_paste() -> Tuple[gradio.Image, gradio.Video]:
    try:
        image = ImageGrab.grabclipboard()
        if isinstance(image, Image.Image):
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            image.save(temp_file.name)
            state_manager.set_item('target_path', temp_file.name)
            return gradio.Image(value=temp_file.name, visible=True), gradio.Video(value=None, visible=False)
    except Exception as e:
        print(f"Error al manejar el portapapeles: {e}")
    return gradio.Image(value=None, visible=False), gradio.Video(value=None, visible=False)


# Función para manejar la carga de archivos en Target
def cargar_archivo(archivo_seleccionado: str) -> Tuple[gradio.Image, gradio.Video]:
    state_manager.set_item('target_path', archivo_seleccionado)
    if is_image(archivo_seleccionado):
        return gradio.Image(value=archivo_seleccionado, visible=True), gradio.Video(value=None, visible=False)
    if is_video(archivo_seleccionado):
        return gradio.Image(value=None, visible=False), gradio.Video(value=archivo_seleccionado, visible=True)
    return gradio.Image(value=None, visible=False), gradio.Video(value=None, visible=False)