Skip to content
This repository was archived by the owner on Dec 14, 2023. It is now read-only.

Commit fb6beda

Browse files
Use Diffusers UNet for inference
We don't need to use the unet model from this repository. It's only for training as some layers are skipped to compensate for spatial and temporal layer training. Once trained, this can be used with any ModelScope implementation 🙂 .
1 parent 87beeaa commit fb6beda

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

inference.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from uuid import uuid4
66
from utils.lora import inject_inferable_lora
77
import torch
8-
from diffusers import DPMSolverMultistepScheduler, TextToVideoSDPipeline
8+
from diffusers import DPMSolverMultistepScheduler, TextToVideoSDPipeline, UNet3DConditionModel
99
from einops import rearrange
1010
from torch.nn.functional import interpolate
1111

@@ -17,8 +17,10 @@ def initialize_pipeline(model, device="cuda", xformers=False, sdp=False):
1717
with warnings.catch_warnings():
1818
warnings.simplefilter("ignore")
1919

20-
scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(model)
21-
20+
scheduler, tokenizer, text_encoder, vae, _unet = load_primary_models(model)
21+
del _unet #This is a no op
22+
unet = UNet3DConditionModel.from_pretrained(model, subfolder='unet')
23+
2224
pipeline = TextToVideoSDPipeline.from_pretrained(
2325
pretrained_model_name_or_path=model,
2426
scheduler=scheduler,

0 commit comments

Comments
 (0)