mirror of
https://github.com/THUDM/CogVideo.git
synced 2025-12-02 10:32:09 +08:00
957a210a72
- Integrated progress tracking with upscale model loading. - Implemented conditional latent upscaling using the upscale model. - Processed batch video frames using PyTorch and converted them to PIL images. - load_torch_file for params_ema convert weights
29 lines
794 B
Python
29 lines
794 B
Python
import utils
|
|
|
|
import torch
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def load_sd_upscale(ckpt):
|
|
from spandrel import ModelLoader, ImageModelDescriptor # Simulate a step in loading
|
|
|
|
pbar = utils.ProgressBar(1, desc="Loading upscale model")
|
|
sd = utils.load_torch_file(ckpt, device=device)
|
|
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
|
|
sd = utils.state_dict_prefix_replace(sd, {"module.": ""})
|
|
out = ModelLoader().load_from_state_dict(sd).half()
|
|
|
|
pbar.update(1) # Update progress by 1
|
|
return out
|
|
|
|
|
|
def test_load_sd_upscale():
|
|
model = load_sd_upscale("/media/gpt4-pdf-chatbot-langchain/ComfyUI/models/upscale_models/RealESRNet_x4plus.pth")
|
|
|
|
print(model.dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_load_sd_upscale()
|