Extras/ffff_assets/source.py

from typing import List, Optional, Tuple
import gradio
import os
import requests
import tempfile
from PIL import ImageGrab, Image  # Para manejar imágenes desde el portapapeles

from ffff import state_manager, wording
from ffff.common_helper import get_first
from ffff.filesystem import filter_audio_paths, filter_image_paths, has_audio, has_image, is_video, get_file_size
from ffff.vision import get_video_frame, normalize_frame_color
from ffff.uis.core import register_ui_component
from ffff.uis.typing import File

FILE_SIZE_LIMIT = 512 * 1024 * 1024

# Variables globales para almacenar el archivo cargado
SOURCE_FILE: Optional[gradio.File] = None
SOURCE_IMAGE: Optional[gradio.Image] = None
SOURCE_VIDEO: Optional[gradio.Video] = None


def render() -> None:
    global SOURCE_FILE
    global SOURCE_IMAGE
    global SOURCE_VIDEO

    has_source_audio = has_audio(state_manager.get_item('source_paths'))
    has_source_image = has_image(state_manager.get_item('source_paths'))
    has_source_video = is_video(state_manager.get_item('source_paths'))

    SOURCE_FILE = gradio.File(
        file_count='single',
        file_types=['image', 'video'],
        label=wording.get('uis.source_file'),
        value=state_manager.get_item('source_paths') if has_source_audio or has_source_image or has_source_video else None
    )

    source_file_path = state_manager.get_item('source_paths')[0] if state_manager.get_item('source_paths') else None
    source_image_args = {'show_label': False, 'visible': False}
    source_video_args = {'show_label': False, 'visible': False}

    if has_source_image:
        source_image_args['value'] = source_file_path
        source_image_args['visible'] = True

    if has_source_video:
        if get_file_size(source_file_path) > FILE_SIZE_LIMIT:
            preview_vision_frame = normalize_frame_color(get_video_frame(source_file_path))
            source_image_args['value'] = preview_vision_frame
            source_image_args['visible'] = True
        else:
            source_video_args['value'] = source_file_path
            source_video_args['visible'] = True

    SOURCE_IMAGE = gradio.Image(**source_image_args)
    SOURCE_VIDEO = gradio.Video(**source_video_args)

    register_ui_component('source_image', SOURCE_IMAGE)
    register_ui_component('source_video', SOURCE_VIDEO)


def listen() -> None:
    SOURCE_FILE.change(update, inputs=SOURCE_FILE, outputs=[SOURCE_IMAGE, SOURCE_VIDEO])


def update(file: File) -> Tuple[gradio.Image, gradio.Video]:
    if file and has_image([file.name]):
        state_manager.set_item('source_paths', [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('source_paths', [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('source_paths')
    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


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('source_paths', [downloaded_file_path])
        if has_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)


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('source_paths', [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)


def cargar_archivo(archivo_seleccionado: str) -> Tuple[gradio.Image, gradio.Video]:
    state_manager.set_item('source_paths', [archivo_seleccionado])
    if has_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)