Spaces:
Running
on
Zero
Running
on
Zero
Commit
·
7a0913a
1
Parent(s):
da57cc8
Fixed model loader
Browse files
app.py
CHANGED
@@ -2,7 +2,7 @@ import spaces
|
|
2 |
import torch
|
3 |
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler, WanTransformer3DModel, AutoModel, DiffusionPipeline
|
4 |
from diffusers.utils import export_to_video
|
5 |
-
from transformers import CLIPVisionModel, UMT5EncoderModel,
|
6 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
7 |
import tempfile
|
8 |
import re
|
@@ -32,8 +32,8 @@ try:
|
|
32 |
# Load ALL components needed for the pipeline from the base model repo
|
33 |
i2v_image_encoder = CLIPVisionModel.from_pretrained(I2V_BASE_MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
|
34 |
i2v_vae = AutoencoderKLWan.from_pretrained(I2V_BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
|
35 |
-
i2v_text_encoder =
|
36 |
-
i2v_tokenizer =
|
37 |
i2v_image_processor = CLIPImageProcessor.from_pretrained(I2V_BASE_MODEL_ID, subfolder="image_processor")
|
38 |
|
39 |
# Create scheduler with custom flow_shift
|
@@ -43,7 +43,8 @@ try:
|
|
43 |
|
44 |
# Load the main transformer from the repo and filename
|
45 |
i2v_transformer = WanTransformer3DModel.from_single_file(
|
46 |
-
|
|
|
47 |
torch_dtype=torch.bfloat16
|
48 |
)
|
49 |
|
|
|
2 |
import torch
|
3 |
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler, WanTransformer3DModel, AutoModel, DiffusionPipeline
|
4 |
from diffusers.utils import export_to_video
|
5 |
+
from transformers import CLIPVisionModel, UMT5EncoderModel, CLIPTextModel, CLIPImageProcessor
|
6 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
7 |
import tempfile
|
8 |
import re
|
|
|
32 |
# Load ALL components needed for the pipeline from the base model repo
|
33 |
i2v_image_encoder = CLIPVisionModel.from_pretrained(I2V_BASE_MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
|
34 |
i2v_vae = AutoencoderKLWan.from_pretrained(I2V_BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
|
35 |
+
i2v_text_encoder = UMT5EncoderModel.from_pretrained(I2V_BASE_MODEL_ID, subfolder="text_encoder", torch_dtype=torch.bfloat16)
|
36 |
+
i2v_tokenizer = AutoTokenizer.from_pretrained(I2V_BASE_MODEL_ID, subfolder="tokenizer")
|
37 |
i2v_image_processor = CLIPImageProcessor.from_pretrained(I2V_BASE_MODEL_ID, subfolder="image_processor")
|
38 |
|
39 |
# Create scheduler with custom flow_shift
|
|
|
43 |
|
44 |
# Load the main transformer from the repo and filename
|
45 |
i2v_transformer = WanTransformer3DModel.from_single_file(
|
46 |
+
I2V_FUSIONX_REPO_ID,
|
47 |
+
filename=I2V_FUSIONX_FILENAME,
|
48 |
torch_dtype=torch.bfloat16
|
49 |
)
|
50 |
|