
    
3j                         S SK Jr  S SKrS SKJr  S SKrS SKrS SKr	S SK
r
S SKJrJr  SSKJr  SSKJr  SSKJr   " S S	5      r " S
 S5      r\\/rg)    )PathN)interpolatorutil   )ModelNotFound)log)get_model_pathc                   d    \ rS rSrSr\S\\   4S j5       r\	S 5       r
SrSrSrS	rS
\4S jrSrg)MTB_LoadFilmModel   zHLoads a FILM model

[DEPRECATED] Use ComfyUI-FrameInterpolation instead
returnc                      [        S5      R                  5       n U  Vs/ s H  oR                  S;   d  M  UPM     sn$ s  snf )NFILM)z.onnxz.pth)r	   iterdirsuffix)models_pathsxs     P/home/wildlama/comfy/ComfyUI/custom_nodes/comfy_mtb/nodes/image_interpolation.py
get_modelsMTB_LoadFilmModel.get_models   s7    %f-557'I<a887H+H<IIIs   >>c                     SS/ SQSS0400$ )Nrequired
film_model)L1StyleVGGdefaultr    clss    r   INPUT_TYPESMTB_LoadFilmModel.INPUT_TYPES   s)     *(
 	
    
FILM_MODEL
load_modelmtb/frame iterpolationTr   c                    [        SU5      nU(       a  UR                  5       (       d  [        SU S35      eUS-  R                  5       (       d  US-  nUR                  5       (       d)  [        R                  " SU S35        [        SU S35      e[        R                  " SU 35        [        R                  " UR                  5       S 5      4$ )	Nr   zFILM ()zsaved_model.pbsaved_modelzModel z does not existzLoading model )
r	   existsr   r   error
ValueErrorinfor   Interpolatoras_posix)selfr   
model_paths      r   r&   MTB_LoadFilmModel.load_model,   s    #FJ7
!2!2!4!4&A 677--5577#m3J  ""IIzl/:;vj\ABB>*./))**=*=*?FHHr#   r   N)__name__
__module____qualname____firstlineno____doc__staticmethodlistr   r   classmethodr!   RETURN_TYPESFUNCTIONCATEGORY
DEPRECATEDstrr&   __static_attributes__r   r#   r   r   r      sa    
 JT
 J J
 
 
 #LH'HJIS Ir#   r   c                   v    \ rS rSrSr\S 5       rSrSrSr	Sr
S\R                  S	\S
\R                  4S jrSrg)MTB_FilmInterpolation=   zoGoogle Research FILM frame interpolation for large motion

[DEPRECATED] Use ComfyUI-FrameInterpolation instead
c                     SSSSSSS.4SS	.0$ )
Nr   IMAGEINTr      2   )r   minmaxr$   )imagesinterpolater   r   r   s    r   r!   !MTB_FilmInterpolation.INPUT_TYPESC   s+     $ %1Qr'JK-
 	
r#   rF   do_interpolationr'   TrM   rN   r   c                 (   UR                  S5      nUS:X  a  U4$ [        R                  R                  S5      n[	        U5      (       d  [
        R                  " S5        O[
        R                  " SU 35        US-
  SU-  S-
  -  n[
        R                  " SU S35        [        U5       Vs/ s H  oqU   PM	     nn/ n	[        R                  R                  U5      n
[        R                  " XU5       Ho  nU	R                  [        U[         R"                  5      (       a  [$        R&                  " U5      OU5        [(        R*                  " 5         U
R-                  S5        Mq     [$        R.                  " U	 Vs/ s H  oR1                  S5      PM     snSS	9n	[
        R                  " S
[	        U	5       S35        [
        R                  " SU	R2                   35        [
        R                  " SU	R4                   35        U	4$ s  snf s  snf )Nr   GPUzHTensorflow GPU not available, falling back to CPU this will be very slowz Tensorflow GPU available, using rI   r   zWill interpolate into z frames)dimz
Returning z tensorszOutput shape zOutput type )sizetfconfiglist_physical_deviceslenr   warningdebugrangecomfyutilsProgressBarr   #interpolate_recursively_from_memoryappend
isinstancenpndarraytorch
from_numpymodel_management)throw_exception_if_processing_interruptedupdatecat	unsqueezeshapedtype)r1   rM   rN   r   navailable_gpus
num_framesi	in_framesout_tensorspbarframetenss                r   rP   &MTB_FilmInterpolation.do_interpolationR   s    KKN69 88?>""KKZ II88HIJ!ek 2Q 67
		*:,g>?(-a11AY	1{{&&z2==J
E eRZZ00   '
 FFHKKN
 ii+67;4^^A;7Q
 			Js;/09:		M+"3"3!456		L!2!2 345~/ 2" 8s   /H
>Hr   N)r4   r5   r6   r7   r8   r;   r!   r<   r=   r>   r?   rd   Tensorintr   r/   rP   rA   r   r#   r   rC   rC   =   s\    
 
 
 L!H'HJ.. . !--	.r#   rC   )pathlibr   r\   comfy.model_managementrf   comfy.utilsnumpyrb   
tensorflowrU   rd   frame_interpolation.evalr   r   errorsr   r   r]   r	   r   rC   	__nodes__r   r#   r   <module>r      sM      1     7 "  "*I *IZC CL  56	r#   