rahul7star commited on
Commit
5474f17
·
verified ·
1 Parent(s): 5ec1f87

Create app1.py

Browse files
Files changed (1) hide show
  1. app1.py +281 -0
app1.py ADDED
@@ -0,0 +1,281 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
4
+
5
+ # wan2.2-main/gradio_ti2v.py
6
+ import gradio as gr
7
+ import torch
8
+ from huggingface_hub import snapshot_download
9
+ from PIL import Image
10
+ import random
11
+ import numpy as np
12
+ import spaces
13
+
14
+ import wan
15
+ from wan.configs import WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS, SUPPORTED_SIZES
16
+ from wan.utils.utils import cache_video
17
+
18
+ import gc
19
+
20
+ # --- 1. Global Setup and Model Loading ---
21
+
22
+ print("Starting Gradio App for Wan 2.2 TI2V-5B...")
23
+
24
+ # Download model snapshots from Hugging Face Hub
25
+ repo_id = "Wan-AI/Wan2.2-TI2V-5B"
26
+ print(f"Downloading/loading checkpoints for {repo_id}...")
27
+ ckpt_dir = snapshot_download(repo_id, local_dir_use_symlinks=False)
28
+ print(f"Using checkpoints from {ckpt_dir}")
29
+
30
+ # Load the model configuration
31
+ TASK_NAME = 'ti2v-5B'
32
+ cfg = WAN_CONFIGS[TASK_NAME]
33
+ FIXED_FPS = 24
34
+ MIN_FRAMES_MODEL = 8
35
+ MAX_FRAMES_MODEL = 121
36
+
37
+ # Dimension calculation constants
38
+ MOD_VALUE = 32
39
+ DEFAULT_H_SLIDER_VALUE = 704
40
+ DEFAULT_W_SLIDER_VALUE = 1280
41
+ NEW_FORMULA_MAX_AREA = 1280.0 * 704.0
42
+
43
+ SLIDER_MIN_H, SLIDER_MAX_H = 128, 1280
44
+ SLIDER_MIN_W, SLIDER_MAX_W = 128, 1280
45
+
46
+ # Instantiate the pipeline in the global scope
47
+ print("Initializing WanTI2V pipeline...")
48
+ device = "cuda" if torch.cuda.is_available() else "cpu"
49
+ device_id = 0 if torch.cuda.is_available() else -1
50
+ pipeline = wan.WanTI2V(
51
+ config=cfg,
52
+ checkpoint_dir=ckpt_dir,
53
+ device_id=device_id,
54
+ rank=0,
55
+ t5_fsdp=False,
56
+ dit_fsdp=False,
57
+ use_sp=False,
58
+ t5_cpu=False,
59
+ init_on_cpu=False,
60
+ convert_model_dtype=True,
61
+ )
62
+ print("Pipeline initialized and ready.")
63
+
64
+ # --- Helper Functions (from Wan 2.1 Fast demo) ---
65
+ def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
66
+ min_slider_h, max_slider_h,
67
+ min_slider_w, max_slider_w,
68
+ default_h, default_w):
69
+ orig_w, orig_h = pil_image.size
70
+ if orig_w <= 0 or orig_h <= 0:
71
+ return default_h, default_w
72
+
73
+ aspect_ratio = orig_h / orig_w
74
+
75
+ calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
76
+ calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
77
+
78
+ calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
79
+ calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
80
+
81
+ new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
82
+ new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
83
+
84
+ return new_h, new_w
85
+
86
+ def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
87
+ """
88
+ Handle image upload and calculate appropriate dimensions for video generation.
89
+
90
+ Args:
91
+ uploaded_pil_image: The uploaded image (PIL Image or numpy array)
92
+ current_h_val: Current height slider value
93
+ current_w_val: Current width slider value
94
+
95
+ Returns:
96
+ Tuple of gr.update objects for height and width sliders
97
+ """
98
+ if uploaded_pil_image is None:
99
+ return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
100
+ try:
101
+ # Convert numpy array to PIL Image if needed
102
+ if hasattr(uploaded_pil_image, 'shape'): # numpy array
103
+ pil_image = Image.fromarray(uploaded_pil_image).convert("RGB")
104
+ else: # already PIL Image
105
+ pil_image = uploaded_pil_image
106
+
107
+ new_h, new_w = _calculate_new_dimensions_wan(
108
+ pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
109
+ SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
110
+ DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
111
+ )
112
+ return gr.update(value=new_h), gr.update(value=new_w)
113
+ except Exception as e:
114
+ gr.Warning("Error attempting to calculate new dimensions")
115
+ return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
116
+
117
+ def get_duration(image,
118
+ prompt,
119
+ height,
120
+ width,
121
+ duration_seconds,
122
+ sampling_steps,
123
+ guide_scale,
124
+ shift,
125
+ seed,
126
+ progress):
127
+ """Calculate dynamic GPU duration based on parameters."""
128
+ if duration_seconds >= 3:
129
+ return 220
130
+ elif sampling_steps > 35 and duration_seconds >= 2:
131
+ return 180
132
+ elif sampling_steps < 35 or duration_seconds < 2:
133
+ return 105
134
+ else:
135
+ return 90
136
+
137
+ # --- 2. Gradio Inference Function ---
138
+ @spaces.GPU(duration=get_duration)
139
+ def generate_video(
140
+ image,
141
+ prompt,
142
+ height,
143
+ width,
144
+ duration_seconds,
145
+ sampling_steps=38,
146
+ guide_scale=cfg.sample_guide_scale,
147
+ shift=cfg.sample_shift,
148
+ seed=42,
149
+ progress=gr.Progress(track_tqdm=True)
150
+ ):
151
+ """
152
+ Generate a video from text prompt and optional image using the Wan 2.2 TI2V model.
153
+
154
+ Args:
155
+ image: Optional input image (numpy array) for image-to-video generation
156
+ prompt: Text prompt describing the desired video
157
+ height: Target video height in pixels
158
+ width: Target video width in pixels
159
+ duration_seconds: Desired video duration in seconds
160
+ sampling_steps: Number of denoising steps for video generation
161
+ guide_scale: Guidance scale for classifier-free guidance
162
+ shift: Sample shift parameter for the model
163
+ seed: Random seed for reproducibility (-1 for random)
164
+ progress: Gradio progress tracker
165
+
166
+ Returns:
167
+ Path to the generated video file
168
+ """
169
+ if seed == -1:
170
+ seed = random.randint(0, sys.maxsize)
171
+
172
+ # Ensure dimensions are multiples of MOD_VALUE
173
+ target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
174
+ target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
175
+
176
+ input_image = None
177
+ if image is not None:
178
+ input_image = Image.fromarray(image).convert("RGB")
179
+ # Resize image to match target dimensions
180
+ input_image = input_image.resize((target_w, target_h))
181
+
182
+ # Calculate number of frames based on duration
183
+ num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
184
+
185
+ # Create size string for the pipeline
186
+ size_str = f"{target_h}*{target_w}"
187
+
188
+ video_tensor = pipeline.generate(
189
+ input_prompt=prompt,
190
+ img=input_image, # Pass None for T2V, Image for I2V
191
+ size=SIZE_CONFIGS.get(size_str, (target_h, target_w)),
192
+ max_area=MAX_AREA_CONFIGS.get(size_str, target_h * target_w),
193
+ frame_num=num_frames, # Use calculated frames instead of cfg.frame_num
194
+ shift=shift,
195
+ sample_solver='unipc',
196
+ sampling_steps=int(sampling_steps),
197
+ guide_scale=guide_scale,
198
+ seed=seed,
199
+ offload_model=True
200
+ )
201
+
202
+ # Save the video to a temporary file
203
+ video_path = cache_video(
204
+ tensor=video_tensor[None], # Add a batch dimension
205
+ save_file=None, # cache_video will create a temp file
206
+ fps=cfg.sample_fps,
207
+ normalize=True,
208
+ value_range=(-1, 1)
209
+ )
210
+ del video_tensor
211
+ gc.collect()
212
+ return video_path
213
+
214
+
215
+ # --- 3. Gradio Interface ---
216
+ css = ".gradio-container {max-width: 1100px !important; margin: 0 auto} #output_video {height: 500px;} #input_image {height: 500px;}"
217
+
218
+ with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
219
+ gr.Markdown("# Wan 2.2 TI2V 5B")
220
+ gr.Markdown("generate high quality videos using **Wan 2.2 5B Text-Image-to-Video model**,[[model]](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B),[[paper]](https://arxiv.org/abs/2503.20314)")
221
+
222
+ with gr.Row():
223
+ with gr.Column(scale=2):
224
+ image_input = gr.Image(type="numpy", label="Optional (blank = text-to-image)", elem_id="input_image")
225
+ prompt_input = gr.Textbox(label="Prompt", value="A beautiful waterfall in a lush jungle, cinematic.", lines=3)
226
+ duration_input = gr.Slider(
227
+ minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1),
228
+ maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1),
229
+ step=0.1,
230
+ value=2.0,
231
+ label="Duration (seconds)",
232
+ info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
233
+ )
234
+
235
+ with gr.Accordion("Advanced Settings", open=False):
236
+ with gr.Row():
237
+ height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
238
+ width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
239
+ steps_input = gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=38, step=1)
240
+ scale_input = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=cfg.sample_guide_scale, step=0.1)
241
+ shift_input = gr.Slider(label="Sample Shift", minimum=1.0, maximum=20.0, value=cfg.sample_shift, step=0.1)
242
+ seed_input = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
243
+
244
+ with gr.Column(scale=2):
245
+ video_output = gr.Video(label="Generated Video", elem_id="output_video")
246
+ run_button = gr.Button("Generate Video", variant="primary")
247
+
248
+ # Add image upload handler
249
+ image_input.upload(
250
+ fn=handle_image_upload_for_dims_wan,
251
+ inputs=[image_input, height_input, width_input],
252
+ outputs=[height_input, width_input]
253
+ )
254
+
255
+ image_input.clear(
256
+ fn=handle_image_upload_for_dims_wan,
257
+ inputs=[image_input, height_input, width_input],
258
+ outputs=[height_input, width_input]
259
+ )
260
+
261
+ example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
262
+ gr.Examples(
263
+ examples=[
264
+ [example_image_path, "The cat removes the glasses from its eyes.", 1088, 800, 1.5],
265
+ [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", 704, 1280, 2.0],
266
+ [None, "Drone footage flying over a futuristic city with flying cars.", 704, 1280, 2.0],
267
+ ],
268
+ inputs=[image_input, prompt_input, height_input, width_input, duration_input],
269
+ outputs=video_output,
270
+ fn=generate_video,
271
+ cache_examples="lazy",
272
+ )
273
+
274
+ run_button.click(
275
+ fn=generate_video,
276
+ inputs=[image_input, prompt_input, height_input, width_input, duration_input, steps_input, scale_input, shift_input, seed_input],
277
+ outputs=video_output
278
+ )
279
+
280
+ if __name__ == "__main__":
281
+ demo.launch(mcp_server=True)