ovi054 commited on
Commit
812e69e
·
verified ·
1 Parent(s): 4b3bba7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from diffusers import UniPCMultistepScheduler
3
+ from diffusers import WanPipeline, AutoencoderKLWan # Use Wan-specific VAE
4
+ from diffusers.models import UNetSpatioTemporalConditionModel
5
+ from transformers import T5EncoderModel, T5Tokenizer
6
+
7
+ from PIL import Image
8
+ import numpy as np
9
+
10
+ model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
11
+ vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
12
+ pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
13
+ flow_shift = 5.0 # 5.0 for 720P, 3.0 for 480P
14
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
15
+ pipe.to("cuda")
16
+
17
+ @spaces.GPU()
18
+ def generate(prompt):
19
+ output = pipe(
20
+ prompt=prompt,
21
+ # negative_prompt=negative_prompt,
22
+ height=720,
23
+ width=1280,
24
+ num_frames=1,
25
+ num_inference_steps=28,
26
+ guidance_scale=5.0,
27
+ )
28
+ image = output.frames[0][0]
29
+ image = (image * 255).astype(np.uint8)
30
+ return Image.fromarray(image)
31
+
32
+ iface = gr.Interface(
33
+ fn=generate,
34
+ inputs=[
35
+ gr.Textbox(label="Input prompt"),
36
+ # gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024),
37
+ # gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=1024),
38
+ # gr.Textbox(label="Lora ID", placeholder="Optional"),
39
+ # gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Lora Scale", value=1)
40
+ ],
41
+ outputs=gr.Image(label="output"),
42
+ )
43
+
44
+ iface.launch()