From 6c7b5972dc6e6b7128d6238bdbf6cc7fd56af2a4 Mon Sep 17 00:00:00 2001 From: cellzero Date: Mon, 20 Jan 2025 20:21:28 +0800 Subject: [PATCH] Fix autocast device type --- comfyui/comfyui_nodes.py | 6 +++++- predict_i2v.py | 2 +- predict_i2v_80g.py | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index e775c5f..4d4fb89 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -549,7 +549,7 @@ class Ruyi_I2VSampler: self.initialize_plugins(pipeline, plugin_parameters, steps) # Inference - with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype): + with torch.no_grad(), torch.autocast(device_type=device.type, dtype=pipeline.transformer.dtype): video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) @@ -583,6 +583,10 @@ class Ruyi_I2VSampler: for _lora_path, _lora_weight in zip(ruyi_model.get("loras", []), ruyi_model.get("strength_model", [])): pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight) + + # Log information + if pipeline.transformer.tea_cache.enabled: + print("TeaCache cached steps:", pipeline.transformer.tea_cache.skip_count) return (videos,) diff --git a/predict_i2v.py b/predict_i2v.py index c276af0..f5dfe9d 100644 --- a/predict_i2v.py +++ b/predict_i2v.py @@ -270,7 +270,7 @@ pipeline.transformer.tea_cache.initialize(tea_cache_enabled, tea_cache_threshold pipeline.transformer.enhance_a_video.initialize(enhance_a_video_enabled, enhance_a_video_weight, enhance_a_video_skip_start_steps, enhance_a_video_skip_end_steps, steps) # Generate video -with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype): +with torch.no_grad(), torch.autocast(device_type=device.type, dtype=pipeline.transformer.dtype): video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) diff --git a/predict_i2v_80g.py b/predict_i2v_80g.py index 182b7fa..aa312f4 100644 --- a/predict_i2v_80g.py +++ b/predict_i2v_80g.py @@ -270,7 +270,7 @@ pipeline.transformer.tea_cache.initialize(tea_cache_enabled, tea_cache_threshold pipeline.transformer.enhance_a_video.initialize(enhance_a_video_enabled, enhance_a_video_weight, enhance_a_video_skip_start_steps, enhance_a_video_skip_end_steps, steps) # Generate video -with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype): +with torch.no_grad(), torch.autocast(device_type=device.type, dtype=pipeline.transformer.dtype): video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width))