add google-veo-3.1

Signed-off-by: vladmandic <mandic00@live.com>

modules/video_models/google_veo.py

@@ -0,0 +1,145 @@
+import io
+import os
+import time
+
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+from PIL import Image
+from installer import install, reload, log
+
+
+image_size_buckets = {
+    '720p': 1280*720,
+    '1080p': 1920*1080,
+}
+aspect_ratios_buckets = {
+    '1:1': 1/1,
+    '2:3': 2/3,
+    '3:2': 3/2,
+    '4:3': 4/3,
+    '3:4': 3/4,
+    '4:5': 4/5,
+    '5:4': 5/4,
+    '16:9': 16/9,
+    '9:16': 9/16,
+    '21:9': 21/9,
+    '9:21': 9/21,
+}
+
+
+def google_requirements():
+    install('google-genai==1.52.0')
+    install('pydantic==2.11.7', ignore=True, quiet=True)
+    reload('pydantic', '2.11.7')
+
+
+def get_size_buckets(width: int, height: int) -> str:
+    aspect_ratio = width / height
+    closest_aspect_ratio = min(aspect_ratios_buckets.items(), key=lambda x: abs(x[1] - aspect_ratio))[0]
+    pixel_count = width * height
+    closest_size = min(image_size_buckets.items(), key=lambda x: abs(x[1] - pixel_count))[0]
+    closest_aspect_ratio = min(aspect_ratios_buckets.items(), key=lambda x: abs(x[1] - aspect_ratio))[0]
+    return closest_size, closest_aspect_ratio
+
+
+class GoogleVeoVideoPipeline():
+    def __init__(self, model_name: str):
+        self.model = model_name
+        self.client = None
+        self.config = None
+        google_requirements()
+        log.debug(f'Load model: type=GoogleVeo model="{model_name}"')
+
+    def txt2vid(self, prompt):
+        return self.client.models.generate_videos(
+            model=self.model,
+            prompt=prompt,
+            config=self.config,
+        )
+
+    def img2vid(self, prompt, image):
+        from google import genai
+        image_bytes = io.BytesIO()
+        image.save(image_bytes, format='JPEG')
+        return self.client.models.generate_videos(
+            model=self.model,
+            prompt=prompt,
+            config=self.config,
+            image=genai.types.Image(image_bytes=image_bytes.getvalue(), mime_type='image/jpeg'),
+        )
+
+    def __call__(self, prompt: list[str], width: int, height: int, image: Image.Image = None, num_frames: int = 4*24):
+        from google import genai
+
+        if isinstance(prompt, list) and len(prompt) > 0:
+            prompt = prompt[0]
+        if self.client is None:
+            api_key = os.getenv("GOOGLE_API_KEY", None)
+            if api_key is None:
+                log.error(f'Cloud: model="{self.model}" GOOGLE_API_KEY environment variable not set')
+                return None
+            self.client = genai.Client(api_key=api_key, vertexai=False)
+
+        resolution, aspect_ratio = get_size_buckets(width, height)
+        duration = num_frames // 24
+        if duration < 4:
+            duration = 4
+        if duration > 8:
+            duration = 8
+        self.config=genai.types.GenerateVideosConfig(
+            # seed=42,
+            # fps=24,
+            duration_seconds=duration,
+            aspect_ratio=aspect_ratio,
+            resolution=resolution,
+            # person_generation='ALLOW_ALL',
+            # safety_filter_level='BLOCK_NONE',
+            # negative_prompt=None,
+            # enhance_prompt=True,
+            # generate_audio=True,
+        )
+        log.debug(f'Cloud: prompt="{prompt}" size={resolution} ar={aspect_ratio} image={image} model="{self.model}" frames={num_frames} duration={duration}')
+
+        operation = None
+        try:
+            if image is not None:
+                operation = self.img2vid(prompt, image)
+            else:
+                operation = self.txt2vid(prompt)
+            while not operation.done:
+                log.debug(f"Cloud processing: {operation}")
+                time.sleep(10)
+                operation = self.client.operations.get(operation)
+        except Exception as e:
+            log.error(f'Cloud video: model="{self.model}" {operation} {e}')
+            return None
+
+        if operation is None or operation.response is None or operation.response.generated_videos is None or len(operation.response.generated_videos) == 0:
+            log.error(f'Cloud video: model="{self.model}" no response {operation}')
+            return None
+        try:
+            response: genai.types.GeneratedVideo = operation.response.generated_videos[0]
+            self.client.files.download(file=response.video)
+            video_bytes = response.video.video_bytes
+            return { 'bytes': video_bytes, 'images': [] }
+        except Exception as e:
+            log.error(f'Cloud download: model="{self.model}" {e}')
+            return None
+
+
+def load_veo(model_name): # pylint: disable=unused-argument
+    pipe = GoogleVeoVideoPipeline(model_name = model_name)
+    return pipe
+
+
+if __name__ == "__main__":
+    from installer import setup_logging
+    setup_logging()
+    log.info('test')
+    model = GoogleVeoVideoPipeline('veo-3.1-generate-preview')
+    img = Image.open('C:\\Users\\mandi\\OneDrive\\Generative\\Samples\\cartoon.png')
+    vid = model(['A beautiful young woman walking through the fantasy city'], 1280, 720, image=img)
+    if vid is not None:
+        with open("veo.mp4", "wb") as f:
+            f.write(vid['video'])

modules/video_models/models_def.py

@@ -10,6 +10,7 @@ class Model():
     name: str
     url: str = ''
     repo: str = None
+    custom: str = None
     repo_cls: classmethod = None
     repo_revision: str = None
     dit: str = None
@@ -480,6 +481,22 @@ try:
                 te_cls=getattr(transformers, 'Qwen2_5_VLForConditionalGeneration', None),
                 dit_cls=getattr(diffusers, 'Kandinsky5Transformer3DModel', None)),
         ],
+        'Google Veo': [
+            Model(name='Google Veo 3.1 T2V',
+                url='https://gemini.google/overview/video-generation/',
+                repo='veo-3.1-generate-preview',
+                custom='GoogleVeoVideoPipeline',
+                repo_cls=None,
+                te_cls=None,
+                dit_cls=None),
+            Model(name='Google Veo 3.1 I2V',
+                url='https://gemini.google/overview/video-generation/',
+                repo='veo-3.1-generate-preview',
+                custom='GoogleVeoVideoPipeline',
+                repo_cls=None,
+                te_cls=None,
+                dit_cls=None),
+        ],
     }
     t1 = time.time()
     errors = 0
@@ -488,11 +505,14 @@ try:
         for m in model:
             if m.name == 'None':
                 continue
+            """
             if (m.repo_cls is None) or (m.dit_cls is None) or (m.te_cls is None):
                 log.error(f'Video: pipeline="{m.name}" not available')
                 errors += 1
             else:
                 total += 1
+            """
+            total += 1
     log.info(f'Networks: type="video" engines={len(models)} models={total} errors={errors} time={t1 - t0:.2f}')
 except Exception as e:
     models = {}

modules/video_models/video_cache.py

@@ -3,7 +3,7 @@ from modules import shared
 
 
 def apply_teacache_patch(cls):
-    if shared.opts.teacache_enabled:
+    if shared.opts.teacache_enabled and cls is not None:
         from modules import teacache
         shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={cls.__name__}')
         if cls.__name__ == 'LTXVideoTransformer3DModel':

modules/video_models/video_load.py

@@ -20,15 +20,26 @@ def _loader(component):
 loaded_model = None
 
 
+def load_custom(model_name: str):
+    shared.log.debug(f'Video load: module=pipe repo="{model_name}" cls=Custom')
+    if 'veo-3.1' in model_name:
+        from modules.video_models.google_veo import load_veo
+        pipe = load_veo(model_name)
+        return pipe
+    return None
+
+
 def load_model(selected: models_def.Model):
-    if selected is None or selected.te_cls is None or selected.dit_cls is None:
+    if selected is None or selected.repo is None:
         return ''
     global loaded_model # pylint: disable=global-statement
     if not shared.sd_loaded:
         loaded_model = None
     if loaded_model == selected.name:
         return ''
-    sd_models.unload_model_weights()
+    if shared.sd_loaded:
+        sd_models.unload_model_weights()
+
     t0 = time.time()
     jobid = shared.state.begin('Load model')
 
@@ -46,89 +57,99 @@ def load_model(selected: models_def.Model):
     kwargs = video_overrides.load_override(selected, **offline_args)
 
     # text encoder
-    try:
-        load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True)
+    if selected.te_cls is not None:
+        try:
+            load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True)
 
-        # loader deduplication of text-encoder models
-        if selected.te_cls.__name__ == 'T5EncoderModel' and shared.opts.te_shared_t5:
-            selected.te = 'Disty0/t5-xxl'
-            selected.te_folder = ''
-            selected.te_revision = None
-        if selected.te_cls.__name__ == 'UMT5EncoderModel' and shared.opts.te_shared_t5:
-            if 'SDNQ' in selected.name:
-                selected.te = 'Disty0/Wan2.2-T2V-A14B-SDNQ-uint4-svd-r32'
-            else:
-                selected.te = 'Wan-AI/Wan2.2-TI2V-5B-Diffusers'
-            selected.te_folder = 'text_encoder'
-            selected.te_revision = None
-        if selected.te_cls.__name__ == 'LlamaModel' and shared.opts.te_shared_t5:
-            selected.te = 'hunyuanvideo-community/HunyuanVideo'
-            selected.te_folder = 'text_encoder'
-            selected.te_revision = None
-        if selected.te_cls.__name__ == 'Qwen2_5_VLForConditionalGeneration' and shared.opts.te_shared_t5:
-            selected.te = 'ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers'
-            selected.te_folder = 'text_encoder'
-            selected.te_revision = None
+            # loader deduplication of text-encoder models
+            if selected.te_cls.__name__ == 'T5EncoderModel' and shared.opts.te_shared_t5:
+                selected.te = 'Disty0/t5-xxl'
+                selected.te_folder = ''
+                selected.te_revision = None
+            if selected.te_cls.__name__ == 'UMT5EncoderModel' and shared.opts.te_shared_t5:
+                if 'SDNQ' in selected.name:
+                    selected.te = 'Disty0/Wan2.2-T2V-A14B-SDNQ-uint4-svd-r32'
+                else:
+                    selected.te = 'Wan-AI/Wan2.2-TI2V-5B-Diffusers'
+                selected.te_folder = 'text_encoder'
+                selected.te_revision = None
+            if selected.te_cls.__name__ == 'LlamaModel' and shared.opts.te_shared_t5:
+                selected.te = 'hunyuanvideo-community/HunyuanVideo'
+                selected.te_folder = 'text_encoder'
+                selected.te_revision = None
+            if selected.te_cls.__name__ == 'Qwen2_5_VLForConditionalGeneration' and shared.opts.te_shared_t5:
+                selected.te = 'ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers'
+                selected.te_folder = 'text_encoder'
+                selected.te_revision = None
 
-        shared.log.debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("transformers")}')
-        kwargs["text_encoder"] = selected.te_cls.from_pretrained(
-            pretrained_model_name_or_path=selected.te or selected.repo,
-            subfolder=selected.te_folder,
-            revision=selected.te_revision or selected.repo_revision,
-            cache_dir=shared.opts.hfcache_dir,
-            **load_args,
-            **quant_args,
-            **offline_args,
-        )
-    except Exception as e:
-        shared.log.error(f'video load: module=te cls={selected.te_cls.__name__} {e}')
-        errors.display(e, 'video')
+            shared.log.debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("transformers")}')
+            kwargs["text_encoder"] = selected.te_cls.from_pretrained(
+                pretrained_model_name_or_path=selected.te or selected.repo,
+                subfolder=selected.te_folder,
+                revision=selected.te_revision or selected.repo_revision,
+                cache_dir=shared.opts.hfcache_dir,
+                **load_args,
+                **quant_args,
+                **offline_args,
+            )
+        except Exception as e:
+            shared.log.error(f'video load: module=te cls={selected.te_cls.__name__} {e}')
+            errors.display(e, 'video')
 
     # transformer
-    try:
-        def load_dit_folder(dit_folder):
-            if dit_folder is not None and dit_folder not in kwargs:
-                # get a new quant arg on every loop to prevent the quant config classes getting entangled
-                load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
-                shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("diffusers")}')
-                kwargs[dit_folder] = selected.dit_cls.from_pretrained(
-                    pretrained_model_name_or_path=selected.dit or selected.repo,
-                    subfolder=dit_folder,
-                    revision=selected.dit_revision or selected.repo_revision,
-                    cache_dir=shared.opts.hfcache_dir,
-                    **load_args,
-                    **quant_args,
-                    **offline_args,
-                )
-            else:
-                shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} loader={_loader("diffusers")} skip')
+    if selected.dit_cls is not None:
+        try:
+            def load_dit_folder(dit_folder):
+                if dit_folder is not None and dit_folder not in kwargs:
+                    # get a new quant arg on every loop to prevent the quant config classes getting entangled
+                    load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
+                    shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("diffusers")}')
+                    kwargs[dit_folder] = selected.dit_cls.from_pretrained(
+                        pretrained_model_name_or_path=selected.dit or selected.repo,
+                        subfolder=dit_folder,
+                        revision=selected.dit_revision or selected.repo_revision,
+                        cache_dir=shared.opts.hfcache_dir,
+                        **load_args,
+                        **quant_args,
+                        **offline_args,
+                    )
+                else:
+                    shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} loader={_loader("diffusers")} skip')
 
-        if selected.dit_folder is None:
-            selected.dit_folder = ['transformer']
-        if isinstance(selected.dit_folder, list) or isinstance(selected.dit_folder, tuple):
-            for dit_folder in selected.dit_folder: # wan a14b has transformer and transformer_2
-                load_dit_folder(dit_folder)
-        else:
-            load_dit_folder(selected.dit_folder)
-    except Exception as e:
-        shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}')
-        errors.display(e, 'video')
+            if selected.dit_folder is None:
+                selected.dit_folder = ['transformer']
+            if isinstance(selected.dit_folder, list) or isinstance(selected.dit_folder, tuple):
+                for dit_folder in selected.dit_folder: # wan a14b has transformer and transformer_2
+                    load_dit_folder(dit_folder)
+            else:
+                load_dit_folder(selected.dit_folder)
+        except Exception as e:
+            shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}')
+            errors.display(e, 'video')
 
     # model
     try:
-        shared.log.debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}')
-        shared.sd_model = selected.repo_cls.from_pretrained(
-            pretrained_model_name_or_path=selected.repo,
-            revision=selected.repo_revision,
-            cache_dir=shared.opts.hfcache_dir,
-            torch_dtype=devices.dtype,
-            **kwargs,
-            **offline_args,
-        )
+        if selected.repo_cls is None:
+            shared.sd_model = load_custom(selected.repo)
+        else:
+            shared.log.debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}')
+            shared.sd_model = selected.repo_cls.from_pretrained(
+                pretrained_model_name_or_path=selected.repo,
+                revision=selected.repo_revision,
+                cache_dir=shared.opts.hfcache_dir,
+                torch_dtype=devices.dtype,
+                **kwargs,
+                **offline_args,
+            )
     except Exception as e:
         shared.log.error(f'video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__} {e}')
         errors.display(e, 'video')
 
+    if shared.sd_model is None:
+        msg = f'Video load: model="{selected.name}" failed'
+        shared.log.error(msg)
+        return msg
+
     t1 = time.time()
     if shared.sd_model.__class__.__name__.startswith("LTX"):
         shared.sd_model.scheduler.config.use_dynamic_shifting = False
@@ -138,7 +159,7 @@ def load_model(selected: models_def.Model):
     sd_models.set_diffuser_options(shared.sd_model, offload=False)
 
     decode, text, image, slicing, tiling, framewise = False, False, False, False, False, False
-    if selected.vae_hijack and hasattr(shared.sd_model.vae, 'decode'):
+    if selected.vae_hijack and hasattr(shared.sd_model, 'vae') and hasattr(shared.sd_model.vae, 'decode'):
         sd_hijack_vae.init_hijack(shared.sd_model)
         decode = True
     if selected.te_hijack and hasattr(shared.sd_model, 'encode_prompt'):

modules/video_models/video_run.py

@@ -109,7 +109,7 @@ def generate(*args, **kwargs):
     orig_sampler_shift = shared.opts.schedulers_shift
     shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift
     shared.opts.data['schedulers_shift'] = sampler_shift
-    if hasattr(shared.sd_model.scheduler, 'config') and hasattr(shared.sd_model.scheduler, 'register_to_config'):
+    if hasattr(shared.sd_model, 'scheduler') and hasattr(shared.sd_model.scheduler, 'config') and hasattr(shared.sd_model.scheduler, 'register_to_config'):
         if hasattr(shared.sd_model.scheduler.config, 'use_dynamic_shifting'):
             shared.sd_model.scheduler.config.use_dynamic_shifting = dynamic_shift
             shared.sd_model.scheduler.register_to_config(use_dynamic_shifting = dynamic_shift)
@@ -146,15 +146,18 @@ def generate(*args, **kwargs):
     # done
     if err:
         return video_utils.queue_err(err)
-    if processed is None or len(processed.images) == 0:
+    if processed is None or (len(processed.images) == 0 and processed.bytes is None):
         return video_utils.queue_err('processing failed')
     shared.log.info(f'Video: name="{selected.name}" cls={shared.sd_model.__class__.__name__} frames={len(processed.images)} time={t1-t0:.2f}')
 
-    # video_file = images.save_video(p, filename=None, images=processed.images, video_type=video_type, duration=video_duration, loop=video_loop, pad=video_pad, interpolate=video_interpolate) # legacy video save from list of images
-    pixels = video_save.images_to_tensor(processed.images)
+    if hasattr(processed, 'images') and processed.images is not None:
+        pixels = video_save.images_to_tensor(processed.images)
+    else:
+        pixels = None
     _num_frames, video_file = video_save.save_video(
         p=p,
         pixels=pixels,
+        binary=processed.bytes,
         mp4_fps=mp4_fps,
         mp4_codec=mp4_codec,
         mp4_opt=mp4_opt,

modules/video_models/video_save.py

@@ -107,7 +107,8 @@ def atomic_save_video(filename, tensor:torch.Tensor, fps:float=24, codec:str='li
 
 def save_video(
         p:processing.StableDiffusionProcessingVideo,
-        pixels:torch.Tensor,
+        pixels:torch.Tensor=None,
+        binary:bytes=None,
         mp4_fps:int=24,
         mp4_codec:str='libx264',
         mp4_opt:str='',
@@ -121,6 +122,23 @@ def save_video(
         pbar=None, # progress bar for video
     ):
     output_video = None
+
+    if binary is not None:
+        output_filename = get_video_filename(p)
+        output_video = f'{output_filename}.{mp4_ext}'
+        try:
+            try:
+                with open(output_video, 'wb') as f:
+                    f.write(binary)
+                shared.log.info(f'Video output: file="{output_video}" size={len(binary)}')
+                shared.state.outputs(output_video)
+            except Exception as e:
+                shared.log.error(f'Video output: file="{output_video}" {e}')
+        except Exception as e:
+            shared.log.error(f'Video output: file="{output_video}" write error {e}')
+            errors.display(e, 'video')
+        return 0, output_video
+
     if pixels is None:
         return 0, output_video
     if not torch.is_tensor(pixels):

modules/video_models/video_vae.py

@@ -9,6 +9,8 @@ vae_type = None
 def set_vae_params(p):
     global vae_type # pylint: disable=global-statement
     vae_type = p.vae_type
+    if not hasattr(shared.sd_model, 'vae'):
+        return
     if hasattr(shared.sd_model.vae, 'enable_slicing'):
         shared.sd_model.vae.enable_slicing()
     if p.frames > p.vae_tile_frames: