Fabrice-TIERCELIN commited on
Commit
90c61c1
·
verified ·
1 Parent(s): c52a7ea

js="document.getElementById('output_id')?.scrollIntoView({behavior: 'smooth', block: 'start', inline: 'nearest'})"

Browse files
Files changed (1) hide show
  1. app.py +677 -672
app.py CHANGED
@@ -1,673 +1,678 @@
1
- import os
2
- # PyTorch 2.8 (temporary hack)
3
- os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
4
-
5
- # --- 1. Model Download and Setup (Diffusers Backend) ---
6
- try:
7
- import spaces
8
- except:
9
- class spaces():
10
- def GPU(*args, **kwargs):
11
- def decorator(function):
12
- return lambda *dummy_args, **dummy_kwargs: function(*dummy_args, **dummy_kwargs)
13
- return decorator
14
-
15
- import torch
16
- from diffusers import FlowMatchEulerDiscreteScheduler
17
- from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
18
- from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
19
- from diffusers.utils.export_utils import export_to_video
20
- import gradio as gr
21
- import imageio_ffmpeg
22
- import tempfile
23
- import shutil
24
- import subprocess
25
- import time
26
- from datetime import datetime
27
- import numpy as np
28
- from PIL import Image
29
- import random
30
- import math
31
- import traceback
32
- import gc
33
- from gradio_client import Client, handle_file # Import for API call
34
- import zipfile
35
-
36
- # Import optimization and access compiled artifacts
37
- import optimization
38
-
39
- # Import the optimization function from the separate file
40
- from optimization import optimize_pipeline_
41
-
42
- # --- Constants and Model Loading ---
43
- MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
44
-
45
- # --- NEW: Flexible Dimension Constants ---
46
- MAX_DIMENSION = 832
47
- MIN_DIMENSION = 480
48
- DIMENSION_MULTIPLE = 16
49
- SQUARE_SIZE = 480
50
-
51
- MAX_SEED = np.iinfo(np.int32).max
52
-
53
- FIXED_FPS = 24
54
- MIN_FRAMES_MODEL = 8
55
- MAX_FRAMES_MODEL = 81
56
-
57
- MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
58
- MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
59
-
60
- input_image_debug_value = [None]
61
- end_image_debug_value = [None]
62
- prompt_debug_value = [None]
63
- total_second_length_debug_value = [None]
64
- resolution_debug_value = [None]
65
- factor_debug_value = [None]
66
- allocation_time_debug_value = [None]
67
-
68
- default_negative_prompt = "Vibrant colors, overexposure, static, blurred details, subtitles, error, style, artwork, painting, image, still, overall gray, worst quality, low quality, JPEG compression residue, ugly, mutilated, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, malformed limbs, fused fingers, still image, cluttered background, three legs, many people in the background, walking backwards, overexposure, jumpcut, crossfader, "
69
-
70
- transformer = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
71
- subfolder='transformer',
72
- torch_dtype=torch.bfloat16,
73
- device_map='cuda',
74
- )
75
-
76
- transformer_2 = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
77
- subfolder='transformer_2',
78
- torch_dtype=torch.bfloat16,
79
- device_map='cuda',
80
- )
81
-
82
- pipe = WanImageToVideoPipeline.from_pretrained(
83
- MODEL_ID,
84
- transformer = transformer,
85
- transformer_2 = transformer_2,
86
- torch_dtype=torch.bfloat16,
87
- )
88
- pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=8.0)
89
- pipe.to('cuda')
90
-
91
- for i in range(3):
92
- gc.collect()
93
- torch.cuda.synchronize()
94
- torch.cuda.empty_cache()
95
-
96
- optimize_pipeline_(pipe,
97
- image=Image.new('RGB', (MAX_DIMENSION, MIN_DIMENSION)),
98
- prompt='prompt',
99
- height=MIN_DIMENSION,
100
- width=MAX_DIMENSION,
101
- num_frames=MAX_FRAMES_MODEL,
102
- )
103
-
104
- def _escape_html(s: str) -> str:
105
- return (s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))
106
-
107
- def _error_to_html(err: BaseException) -> str:
108
- tb = traceback.format_exc()
109
- return (
110
- "<div style='padding:12px;border:1px solid #ff4d4f;background:#fff1f0;color:#a8071a;border-radius:8px;'>"
111
- "<b>Generation failed</b><br/>"
112
- f"<b>{_escape_html(type(err).__name__)}</b>: {_escape_html(str(err))}"
113
- "<details style='margin-top:8px;'>"
114
- "<summary>Show traceback</summary>"
115
- f"<pre style='white-space:pre-wrap;margin-top:8px;'>{_escape_html(tb)}</pre>"
116
- "</details>"
117
- "</div>"
118
- )
119
-
120
- # 20250508 pftq: for saving prompt to mp4 metadata comments
121
- def set_mp4_comments_imageio_ffmpeg(input_file, comments):
122
- try:
123
- # Get the path to the bundled FFmpeg binary from imageio-ffmpeg
124
- ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
125
-
126
- # Check if input file exists
127
- if not os.path.exists(input_file):
128
- #print(f"Error: Input file {input_file} does not exist")
129
- return False
130
-
131
- # Create a temporary file path
132
- temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
133
-
134
- # FFmpeg command using the bundled binary
135
- command = [
136
- ffmpeg_path, # Use imageio-ffmpeg's FFmpeg
137
- '-i', input_file, # input file
138
- '-metadata', f'comment={comments}', # set comment metadata
139
- '-c:v', 'copy', # copy video stream without re-encoding
140
- '-c:a', 'copy', # copy audio stream without re-encoding
141
- '-y', # overwrite output file if it exists
142
- temp_file # temporary output file
143
- ]
144
-
145
- # Run the FFmpeg command
146
- result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
147
-
148
- if result.returncode == 0:
149
- # Replace the original file with the modified one
150
- shutil.move(temp_file, input_file)
151
- #print(f"Successfully added comments to {input_file}")
152
- return True
153
- else:
154
- # Clean up temp file if FFmpeg fails
155
- if os.path.exists(temp_file):
156
- os.remove(temp_file)
157
- #print(f"Error: FFmpeg failed with message:\n{result.stderr}")
158
- return False
159
-
160
- except Exception as e:
161
- # Clean up temp file in case of other errors
162
- if 'temp_file' in locals() and os.path.exists(temp_file):
163
- os.remove(temp_file)
164
- print(f"Error saving prompt to video metadata, ffmpeg may be required: "+str(e))
165
- return False
166
-
167
- # --- 2. Image Processing and Application Logic ---
168
- def generate_end_frame(start_img, gen_prompt, progress=gr.Progress(track_tqdm=True)):
169
- """Calls an external Gradio API to generate an image."""
170
- if start_img is None:
171
- raise gr.Error("Please provide a Start Frame first.")
172
-
173
- hf_token = os.getenv("HF_TOKEN")
174
- if not hf_token:
175
- raise gr.Error("HF_TOKEN not found in environment variables. Please set it in your Space secrets.")
176
-
177
- with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
178
- start_img.save(tmpfile.name)
179
- tmp_path = tmpfile.name
180
-
181
- progress(0.1, desc="Connecting to image generation API...")
182
- client = Client("multimodalart/nano-banana-private")
183
-
184
- progress(0.5, desc=f"Generating with prompt: '{gen_prompt}'...")
185
- try:
186
- result = client.predict(
187
- prompt=gen_prompt,
188
- images=[
189
- {"image": handle_file(tmp_path)}
190
- ],
191
- manual_token=hf_token,
192
- api_name="/unified_image_generator"
193
- )
194
- finally:
195
- os.remove(tmp_path)
196
-
197
- progress(1.0, desc="Done!")
198
- print(result)
199
- return result
200
-
201
- def switch_to_upload_tab():
202
- """Returns a gr.Tabs update to switch to the first tab."""
203
- return gr.Tabs(selected="upload_tab")
204
-
205
-
206
- def process_image_for_video(image: Image.Image, resolution: int) -> Image.Image:
207
- """
208
- Resizes an image based on the following rules for video generation.
209
- """
210
- width, height = image.size
211
-
212
- if resolution < width * height:
213
- scale = ((width * height) / resolution)**(.5)
214
- new_width = width / scale
215
- new_height = height / scale
216
- final_width = int(math.floor(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
217
- final_height = int(math.floor(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
218
-
219
- elif width * height < (MIN_DIMENSION**2):
220
- scale = ((MIN_DIMENSION**2) / (width * height))**(.5)
221
- new_width = width * scale
222
- new_height = height * scale
223
- final_width = int(math.ceil(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
224
- final_height = int(math.ceil(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
225
-
226
- else:
227
- final_width = int(round(width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
228
- final_height = int(round(height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
229
-
230
- return image.resize((final_width, final_height), Image.Resampling.LANCZOS)
231
-
232
- def resize_and_crop_to_match(target_image, reference_image):
233
- """Resizes the target image to match the reference image's dimensions."""
234
- ref_width, ref_height = reference_image.size
235
- return target_image.resize((ref_width, ref_height), Image.Resampling.LANCZOS)
236
-
237
- def crop_to_match(target_image, reference_image):
238
- """Resizes and center-crops the target image to match the reference image's dimensions."""
239
- ref_width, ref_height = reference_image.size
240
- target_width, target_height = target_image.size
241
- scale = max(ref_width / target_width, ref_height / target_height)
242
- new_width, new_height = int(target_width * scale), int(target_height * scale)
243
- resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
244
- left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
245
- return resized.crop((left, top, left + ref_width, top + ref_height))
246
-
247
- def init_view():
248
- return gr.update(interactive = True)
249
-
250
- def output_video_change(output_video):
251
- print('Log output: ' + str(output_video))
252
- return [gr.update(visible = True)] * 2
253
-
254
- def generate_video(
255
- start_image_pil,
256
- end_image_pil,
257
- prompt,
258
- negative_prompt=default_negative_prompt,
259
- resolution=500000,
260
- duration_seconds=2.1,
261
- steps=8,
262
- guidance_scale=1,
263
- guidance_scale_2=1,
264
- seed=42,
265
- randomize_seed=True,
266
- progress=gr.Progress(track_tqdm=True)
267
- ):
268
- start = time.time()
269
- allocation_time = 120
270
- factor = 1
271
-
272
- if input_image_debug_value[0] is not None or end_image_debug_value[0] is not None or prompt_debug_value[0] is not None or total_second_length_debug_value[0] is not None or allocation_time_debug_value[0] is not None or resolution_debug_value[0] is not None or factor_debug_value[0] is not None:
273
- start_image_pil = input_image_debug_value[0]
274
- end_image_pil = end_image_debug_value[0]
275
- prompt = prompt_debug_value[0]
276
- duration_seconds = total_second_length_debug_value[0]
277
- resolution = resolution_debug_value[0]
278
- factor = factor_debug_value[0]
279
- allocation_time = allocation_time_debug_value[0]
280
-
281
- if start_image_pil is None or end_image_pil is None:
282
- raise gr.Error("Please upload both a start and an end image.")
283
-
284
- # Step 1: Process the start image to get our target dimensions based on the new rules.
285
- processed_start_image = process_image_for_video(start_image_pil, resolution)
286
-
287
- # Step 2: Make the end image match the *exact* dimensions of the processed start image.
288
- processed_end_image = resize_and_crop_to_match(end_image_pil, processed_start_image)
289
-
290
- target_height, target_width = processed_start_image.height, processed_start_image.width
291
-
292
- # Handle seed and frame count
293
- current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
294
- num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
295
-
296
- progress(0.2, desc=f"Generating {num_frames} frames at {target_width}x{target_height} (seed: {current_seed})...")
297
-
298
- progress(0.1, desc="Preprocessing images...")
299
- print("Generate a video with the prompt: " + prompt)
300
- output_frames_list = None
301
- caught_error = None
302
- while factor >= 1 and int(allocation_time) > 0:
303
- try:
304
- output_frames_list = generate_video_on_gpu(
305
- start_image_pil,
306
- end_image_pil,
307
- prompt,
308
- negative_prompt,
309
- int(steps),
310
- float(guidance_scale),
311
- float(guidance_scale_2),
312
- progress,
313
- allocation_time,
314
- target_height,
315
- target_width,
316
- current_seed,
317
- (int(((num_frames * factor) - 1) / 4) * 4) + 1,
318
- processed_start_image,
319
- processed_end_image
320
- )
321
- factor = 0
322
- caught_error = None
323
- except BaseException as err:
324
- print("An exception occurred: " + str(err))
325
- caught_error = err
326
- try:
327
- print('e.message: ' + err.message) # No GPU is currently available for you after 60s
328
- except Exception as e2:
329
- print('Failure')
330
- if not str(err).startswith("No GPU is currently available for you after 60s"):
331
- factor -= .003
332
- allocation_time = int(allocation_time) - 1
333
- except:
334
- print("An error occurred")
335
- caught_error = None
336
- if not str(e).startswith("No GPU is currently available for you after 60s"):
337
- factor -= .003
338
- allocation_time = int(allocation_time) - 1
339
-
340
- if caught_error is not None:
341
- return [gr.skip(), gr.skip(), gr.skip(), gr.update(value=_error_to_html(caught_error), visible=True), gr.skip()]
342
-
343
- input_image_debug_value[0] = end_image_debug_value[0] = prompt_debug_value[0] = total_second_length_debug_value[0] = allocation_time_debug_value[0] = factor_debug_value[0] = None
344
-
345
- progress(0.9, desc="Encoding and saving video...")
346
-
347
- video_path = 'wan_' + datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.mp4'
348
-
349
- export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
350
- set_mp4_comments_imageio_ffmpeg(video_path, f"Prompt: {prompt} | Negative Prompt: {negative_prompt}");
351
- print("Video exported: " + video_path)
352
-
353
- progress(1.0, desc="Done!")
354
- end = time.time()
355
- secondes = int(end - start)
356
- minutes = math.floor(secondes / 60)
357
- secondes = secondes - (minutes * 60)
358
- hours = math.floor(minutes / 60)
359
- minutes = minutes - (hours * 60)
360
- information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
361
- "The video been generated in " + \
362
- ((str(hours) + " h, ") if hours != 0 else "") + \
363
- ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
364
- str(secondes) + " sec (including " + str(allocation_time) + " seconds of GPU). " + \
365
- "The video has " + str(int(num_frames * factor)) + " frames. " + \
366
- "The video resolution is " + str(target_width) + \
367
- " pixels large and " + str(target_height) + \
368
- " pixels high, so a resolution of " + f'{target_width * target_height:,}' + " pixels." + \
369
- " Your prompt is saved into the metadata of the video."
370
- return [video_path, gr.update(value = video_path, visible = True, interactive = True), current_seed, gr.update(value = information, visible = True), gr.update(interactive = False)]
371
-
372
- def get_duration(
373
- start_image_pil,
374
- end_image_pil,
375
- prompt,
376
- negative_prompt,
377
- steps,
378
- guidance_scale,
379
- guidance_scale_2,
380
- progress,
381
- allocation_time,
382
- target_height,
383
- target_width,
384
- current_seed,
385
- num_frames,
386
- processed_start_image,
387
- processed_end_image
388
- ):
389
- return allocation_time
390
-
391
- @torch.no_grad()
392
- @spaces.GPU(duration=get_duration)
393
- def generate_video_on_gpu(
394
- start_image_pil,
395
- end_image_pil,
396
- prompt,
397
- negative_prompt,
398
- steps,
399
- guidance_scale,
400
- guidance_scale_2,
401
- progress,
402
- allocation_time,
403
- target_height,
404
- target_width,
405
- current_seed,
406
- num_frames,
407
- processed_start_image,
408
- processed_end_image
409
- ):
410
- """
411
- Generates a video by interpolating between a start and end image, guided by a text prompt,
412
- using the diffusers Wan2.2 pipeline.
413
- """
414
-
415
- output_frames_list = pipe(
416
- image=processed_start_image,
417
- last_image=processed_end_image,
418
- prompt=prompt,
419
- negative_prompt=negative_prompt,
420
- height=target_height,
421
- width=target_width,
422
- num_frames=num_frames,
423
- guidance_scale=guidance_scale,
424
- guidance_scale_2=guidance_scale_2,
425
- num_inference_steps=steps,
426
- generator=torch.Generator(device="cuda").manual_seed(current_seed),
427
- ).frames[0]
428
-
429
- return output_frames_list
430
-
431
- def export_compiled_transformers_to_zip() -> str:
432
- """
433
- Bundle compiled_transformer_1 and compiled_transformer_2 into a zip file and return the file path.
434
- """
435
- ct1 = getattr(optimization, "COMPILED_TRANSFORMER_1", None)
436
- ct2 = getattr(optimization, "COMPILED_TRANSFORMER_2", None)
437
-
438
- if ct1 is None or ct2 is None:
439
- raise gr.Error("Compiled transformers are not available yet (compilation may have failed).")
440
-
441
- payload_1 = ct1.to_serializable_dict()
442
- payload_2 = ct2.to_serializable_dict()
443
-
444
- tmp_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
445
- tmp_zip.close()
446
-
447
- with zipfile.ZipFile(tmp_zip.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
448
- # store with torch.save so users can load easily with torch.load()
449
- buf1 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
450
- buf1.close()
451
- torch.save(payload_1, buf1.name)
452
-
453
- buf2 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
454
- buf2.close()
455
- torch.save(payload_2, buf2.name)
456
-
457
- zf.write(buf1.name, arcname="compiled_transformer_1.pt")
458
- zf.write(buf2.name, arcname="compiled_transformer_2.pt")
459
-
460
- # cleanup intermediate .pt
461
- try:
462
- os.remove(buf1.name)
463
- os.remove(buf2.name)
464
- except:
465
- pass
466
-
467
- return tmp_zip.name
468
-
469
-
470
- # --- 3. Gradio User Interface ---
471
-
472
-
473
-
474
- js = """
475
- function createGradioAnimation() {
476
- window.addEventListener("beforeunload", function(e) {
477
- if (document.getElementById('dummy_button_id') && !document.getElementById('dummy_button_id').disabled) {
478
- var confirmationMessage = 'A process is still running. '
479
- + 'If you leave before saving, your changes will be lost.';
480
-
481
- (e || window.event).returnValue = confirmationMessage;
482
- }
483
- return confirmationMessage;
484
- });
485
- return 'Animation created';
486
- }
487
- """
488
-
489
- # Gradio interface
490
- with gr.Blocks(js=js) as app:
491
- gr.Markdown("# Wan 2.2 First/Last Frame Video Fast")
492
- gr.Markdown("Based on the [Wan 2.2 First/Last Frame workflow](https://www.reddit.com/r/StableDiffusion/comments/1me4306/psa_wan_22_does_first_frame_last_frame_out_of_the/), applied to 🧨 Diffusers + [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning) 8-step LoRA")
493
-
494
- with gr.Row(elem_id="general_items"):
495
- with gr.Column():
496
- with gr.Group(elem_id="group_all"):
497
- with gr.Row():
498
- start_image = gr.Image(type="pil", label="Start Frame", sources=["upload", "clipboard"])
499
- # Capture the Tabs component in a variable and assign IDs to tabs
500
- with gr.Tabs(elem_id="group_tabs") as tabs:
501
- with gr.TabItem("Upload", id="upload_tab"):
502
- end_image = gr.Image(type="pil", label="End Frame", sources=["upload", "clipboard"])
503
- with gr.TabItem("Generate", id="generate_tab"):
504
- generate_5seconds = gr.Button("Generate scene 5 seconds in the future", elem_id="fivesec")
505
- gr.Markdown("Generate a custom end-frame with an edit model like [Nano Banana](https://huggingface.co/spaces/multimodalart/nano-banana) or [Qwen Image Edit](https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast)", elem_id="or_item")
506
- prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images", placeholder="The creature starts to move")
507
-
508
- with gr.Accordion("Advanced Settings", open=False):
509
- duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=2.1, label="Video Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
510
- negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
511
- resolution = gr.Dropdown([
512
- ["400,000 px (working)", 400000],
513
- ["465,920 px (working)", 465920],
514
- ["495,616 px (working)", 495616],
515
- ["500,000 px (working)", 500000],
516
- ["600,000 px (working)", 600000],
517
- ["700,000 px (working)", 700000],
518
- ["800,000 px (working)", 800000],
519
- ["900,000 px (working)", 900000],
520
- ["1,000,000 px (working)", 1000000],
521
- ["1,100,000 px (untested)", 1100000],
522
- ["1,200,000 px (untested)", 1200000],
523
- ["1,300,000 px (untested)", 1300000],
524
- ["1,400,000 px (untested)", 1400000],
525
- ["1,500,000 px (untested)", 1500000]
526
- ], value=465920, label="Resolution (width x height)", info="Less if the image is smaller")
527
- steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
528
- guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - high noise")
529
- guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - low noise")
530
- with gr.Row():
531
- seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
532
- randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
533
-
534
- generate_button = gr.Button("🚀 Generate Video", variant="primary")
535
- dummy_button = gr.Button(elem_id = "dummy_button_id", visible = False, interactive = False)
536
-
537
- with gr.Column():
538
- output_video = gr.Video(label="Generated Video", autoplay = True, loop = True)
539
- download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
540
- video_information = gr.HTML(value = "")
541
-
542
- with gr.Accordion("🔧 Compilation artifacts (advanced)", open=False):
543
- gr.Markdown(
544
- "Télécharge les artefacts compilés AOTInductor générés au démarrage (transformer + transformer_2)."
545
- )
546
- export_btn = gr.Button("📦 Préparer l'archive des transformers compilés")
547
- compiled_download = gr.DownloadButton(label="⬇️ Télécharger compiled_transformers.zip", interactive=False)
548
-
549
- def _build_and_enable_download():
550
- path = export_compiled_transformers_to_zip()
551
- return gr.update(value=path, interactive=True)
552
-
553
- export_btn.click(fn=_build_and_enable_download, inputs=None, outputs=compiled_download)
554
-
555
- # Main video generation button
556
- ui_inputs = [
557
- start_image,
558
- end_image,
559
- prompt,
560
- negative_prompt_input,
561
- resolution,
562
- duration_seconds_input,
563
- steps_slider,
564
- guidance_scale_input,
565
- guidance_scale_2_input,
566
- seed_input,
567
- randomize_seed_checkbox
568
- ]
569
- ui_outputs = [output_video, download_button, seed_input, video_information, dummy_button]
570
-
571
- generate_button.click(fn = init_view, inputs = [], outputs = [dummy_button], queue = False, show_progress = False).success(
572
- fn = generate_video,
573
- inputs = ui_inputs,
574
- outputs = ui_outputs
575
- )
576
-
577
- generate_5seconds.click(
578
- fn=switch_to_upload_tab,
579
- inputs=None,
580
- outputs=[tabs]
581
- ).then(
582
- fn=lambda img: generate_end_frame(img, "this image is a still frame from a movie. generate a new frame with what happens on this scene 5 seconds in the future"),
583
- inputs=[start_image],
584
- outputs=[end_image]
585
- ).success(
586
- fn=generate_video,
587
- inputs=ui_inputs,
588
- outputs=ui_outputs
589
- )
590
-
591
- output_video.change(
592
- fn=output_video_change,
593
- inputs=[output_video],
594
- outputs=[download_button, video_information],
595
- js="document.getElementById('download_btn').click()"
596
- )
597
-
598
- with gr.Row(visible=False):
599
- gr.Examples(
600
- examples=[["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."]],
601
- inputs=[start_image, end_image, prompt],
602
- outputs=ui_outputs,
603
- fn=generate_video,
604
- run_on_click=True,
605
- cache_examples=True,
606
- cache_mode='lazy',
607
- )
608
- prompt_debug=gr.Textbox(label="Prompt Debug")
609
- input_image_debug=gr.Image(type="pil", label="Image Debug")
610
- end_image_debug=gr.Image(type="pil", label="End Image Debug")
611
- total_second_length_debug=gr.Slider(label="Duration Debug", minimum=1, maximum=120, value=5, step=0.1)
612
- resolution_debug = gr.Dropdown([
613
- ["400,000 px", 400000],
614
- ["465,920 px", 465920],
615
- ["495,616 px", 495616],
616
- ["500,000 px", 500000],
617
- ["600,000 px", 600000],
618
- ["700,000 px", 700000],
619
- ["800,000 px", 800000],
620
- ["900,000 px", 900000],
621
- ["1,000,000 px", 1000000],
622
- ["1,100,000 px", 1100000],
623
- ["1,200,000 px", 1200000],
624
- ["1,300,000 px", 1300000],
625
- ["1,400,000 px", 1400000],
626
- ["1,500,000 px", 1500000]
627
- ], value=500000, label="Resolution Debug")
628
- factor_debug=gr.Slider(label="Factor Debug", minimum=1, maximum=100, value=3.2, step=0.1)
629
- allocation_time_debug=gr.Slider(label="Allocation Debug", minimum=1, maximum=60 * 20, value=720, step=1)
630
-
631
- def handle_field_debug_change(
632
- input_image_debug_data,
633
- end_image_debug_data,
634
- prompt_debug_data,
635
- total_second_length_debug_data,
636
- resolution_debug_data,
637
- factor_debug_data,
638
- allocation_time_debug_data
639
- ):
640
- input_image_debug_value[0] = input_image_debug_data
641
- end_image_debug_value[0] = end_image_debug_data
642
- prompt_debug_value[0] = prompt_debug_data
643
- total_second_length_debug_value[0] = total_second_length_debug_data
644
- resolution_debug_value[0] = resolution_debug_data
645
- factor_debug_value[0] = factor_debug_data
646
- allocation_time_debug_value[0] = allocation_time_debug_data
647
- return []
648
-
649
- inputs_debug=[input_image_debug, end_image_debug, prompt_debug, total_second_length_debug, resolution_debug, factor_debug, allocation_time_debug]
650
-
651
- input_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
652
- end_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
653
- prompt_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
654
- total_second_length_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
655
- resolution_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
656
- factor_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
657
- allocation_time_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
658
-
659
- gr.Examples(
660
- label = "Examples from demo",
661
- examples = [
662
- ["poli_tower.png", "tower_takes_off.png", "The man turns around."],
663
- ["ugly_sonic.jpeg", "squatting_sonic.png", "पात्रं क्षेपणास्त्रं चकमाति।"],
664
- ["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."],
665
- ],
666
- inputs = [start_image, end_image, prompt],
667
- outputs = ui_outputs,
668
- fn = generate_video,
669
- cache_examples = False,
670
- )
671
-
672
- if __name__ == "__main__":
 
 
 
 
 
673
  app.launch(mcp_server=True, share=True)
 
1
+ import os
2
+ # PyTorch 2.8 (temporary hack)
3
+ os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
4
+
5
+ # --- 1. Model Download and Setup (Diffusers Backend) ---
6
+ try:
7
+ import spaces
8
+ except:
9
+ class spaces():
10
+ def GPU(*args, **kwargs):
11
+ def decorator(function):
12
+ return lambda *dummy_args, **dummy_kwargs: function(*dummy_args, **dummy_kwargs)
13
+ return decorator
14
+
15
+ import torch
16
+ from diffusers import FlowMatchEulerDiscreteScheduler
17
+ from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
18
+ from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
19
+ from diffusers.utils.export_utils import export_to_video
20
+ import gradio as gr
21
+ import imageio_ffmpeg
22
+ import tempfile
23
+ import shutil
24
+ import subprocess
25
+ import time
26
+ from datetime import datetime
27
+ import numpy as np
28
+ from PIL import Image
29
+ import random
30
+ import math
31
+ import traceback
32
+ import gc
33
+ from gradio_client import Client, handle_file # Import for API call
34
+ import zipfile
35
+
36
+ # Import optimization and access compiled artifacts
37
+ import optimization
38
+
39
+ # Import the optimization function from the separate file
40
+ from optimization import optimize_pipeline_
41
+
42
+ # --- Constants and Model Loading ---
43
+ MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
44
+
45
+ # --- NEW: Flexible Dimension Constants ---
46
+ MAX_DIMENSION = 832
47
+ MIN_DIMENSION = 480
48
+ DIMENSION_MULTIPLE = 16
49
+ SQUARE_SIZE = 480
50
+
51
+ MAX_SEED = np.iinfo(np.int32).max
52
+
53
+ FIXED_FPS = 24
54
+ MIN_FRAMES_MODEL = 8
55
+ MAX_FRAMES_MODEL = 81
56
+
57
+ MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
58
+ MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
59
+
60
+ input_image_debug_value = [None]
61
+ end_image_debug_value = [None]
62
+ prompt_debug_value = [None]
63
+ total_second_length_debug_value = [None]
64
+ resolution_debug_value = [None]
65
+ factor_debug_value = [None]
66
+ allocation_time_debug_value = [None]
67
+
68
+ default_negative_prompt = "Vibrant colors, overexposure, static, blurred details, subtitles, error, style, artwork, painting, image, still, overall gray, worst quality, low quality, JPEG compression residue, ugly, mutilated, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, malformed limbs, fused fingers, still image, cluttered background, three legs, many people in the background, walking backwards, overexposure, jumpcut, crossfader, "
69
+
70
+ transformer = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
71
+ subfolder='transformer',
72
+ torch_dtype=torch.bfloat16,
73
+ device_map='cuda',
74
+ )
75
+
76
+ transformer_2 = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
77
+ subfolder='transformer_2',
78
+ torch_dtype=torch.bfloat16,
79
+ device_map='cuda',
80
+ )
81
+
82
+ pipe = WanImageToVideoPipeline.from_pretrained(
83
+ MODEL_ID,
84
+ transformer = transformer,
85
+ transformer_2 = transformer_2,
86
+ torch_dtype=torch.bfloat16,
87
+ )
88
+ pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=8.0)
89
+ pipe.to('cuda')
90
+
91
+ for i in range(3):
92
+ gc.collect()
93
+ torch.cuda.synchronize()
94
+ torch.cuda.empty_cache()
95
+
96
+ optimize_pipeline_(pipe,
97
+ image=Image.new('RGB', (MAX_DIMENSION, MIN_DIMENSION)),
98
+ prompt='prompt',
99
+ height=MIN_DIMENSION,
100
+ width=MAX_DIMENSION,
101
+ num_frames=MAX_FRAMES_MODEL,
102
+ )
103
+
104
+ def _escape_html(s: str) -> str:
105
+ return (s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))
106
+
107
+ def _error_to_html(err: BaseException) -> str:
108
+ tb = traceback.format_exc()
109
+ return (
110
+ "<div style='padding:12px;border:1px solid #ff4d4f;background:#fff1f0;color:#a8071a;border-radius:8px;'>"
111
+ "<b>Generation failed</b><br/>"
112
+ f"<b>{_escape_html(type(err).__name__)}</b>: {_escape_html(str(err))}"
113
+ "<details style='margin-top:8px;'>"
114
+ "<summary>Show traceback</summary>"
115
+ f"<pre style='white-space:pre-wrap;margin-top:8px;'>{_escape_html(tb)}</pre>"
116
+ "</details>"
117
+ "</div>"
118
+ )
119
+
120
+ # 20250508 pftq: for saving prompt to mp4 metadata comments
121
+ def set_mp4_comments_imageio_ffmpeg(input_file, comments):
122
+ try:
123
+ # Get the path to the bundled FFmpeg binary from imageio-ffmpeg
124
+ ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
125
+
126
+ # Check if input file exists
127
+ if not os.path.exists(input_file):
128
+ #print(f"Error: Input file {input_file} does not exist")
129
+ return False
130
+
131
+ # Create a temporary file path
132
+ temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
133
+
134
+ # FFmpeg command using the bundled binary
135
+ command = [
136
+ ffmpeg_path, # Use imageio-ffmpeg's FFmpeg
137
+ '-i', input_file, # input file
138
+ '-metadata', f'comment={comments}', # set comment metadata
139
+ '-c:v', 'copy', # copy video stream without re-encoding
140
+ '-c:a', 'copy', # copy audio stream without re-encoding
141
+ '-y', # overwrite output file if it exists
142
+ temp_file # temporary output file
143
+ ]
144
+
145
+ # Run the FFmpeg command
146
+ result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
147
+
148
+ if result.returncode == 0:
149
+ # Replace the original file with the modified one
150
+ shutil.move(temp_file, input_file)
151
+ #print(f"Successfully added comments to {input_file}")
152
+ return True
153
+ else:
154
+ # Clean up temp file if FFmpeg fails
155
+ if os.path.exists(temp_file):
156
+ os.remove(temp_file)
157
+ #print(f"Error: FFmpeg failed with message:\n{result.stderr}")
158
+ return False
159
+
160
+ except Exception as e:
161
+ # Clean up temp file in case of other errors
162
+ if 'temp_file' in locals() and os.path.exists(temp_file):
163
+ os.remove(temp_file)
164
+ print(f"Error saving prompt to video metadata, ffmpeg may be required: "+str(e))
165
+ return False
166
+
167
+ # --- 2. Image Processing and Application Logic ---
168
+ def generate_end_frame(start_img, gen_prompt, progress=gr.Progress(track_tqdm=True)):
169
+ """Calls an external Gradio API to generate an image."""
170
+ if start_img is None:
171
+ raise gr.Error("Please provide a Start Frame first.")
172
+
173
+ hf_token = os.getenv("HF_TOKEN")
174
+ if not hf_token:
175
+ raise gr.Error("HF_TOKEN not found in environment variables. Please set it in your Space secrets.")
176
+
177
+ with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
178
+ start_img.save(tmpfile.name)
179
+ tmp_path = tmpfile.name
180
+
181
+ progress(0.1, desc="Connecting to image generation API...")
182
+ client = Client("multimodalart/nano-banana-private")
183
+
184
+ progress(0.5, desc=f"Generating with prompt: '{gen_prompt}'...")
185
+ try:
186
+ result = client.predict(
187
+ prompt=gen_prompt,
188
+ images=[
189
+ {"image": handle_file(tmp_path)}
190
+ ],
191
+ manual_token=hf_token,
192
+ api_name="/unified_image_generator"
193
+ )
194
+ finally:
195
+ os.remove(tmp_path)
196
+
197
+ progress(1.0, desc="Done!")
198
+ print(result)
199
+ return result
200
+
201
+ def switch_to_upload_tab():
202
+ """Returns a gr.Tabs update to switch to the first tab."""
203
+ return gr.Tabs(selected="upload_tab")
204
+
205
+
206
+ def process_image_for_video(image: Image.Image, resolution: int) -> Image.Image:
207
+ """
208
+ Resizes an image based on the following rules for video generation.
209
+ """
210
+ width, height = image.size
211
+
212
+ if resolution < width * height:
213
+ scale = ((width * height) / resolution)**(.5)
214
+ new_width = width / scale
215
+ new_height = height / scale
216
+ final_width = int(math.floor(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
217
+ final_height = int(math.floor(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
218
+
219
+ elif width * height < (MIN_DIMENSION**2):
220
+ scale = ((MIN_DIMENSION**2) / (width * height))**(.5)
221
+ new_width = width * scale
222
+ new_height = height * scale
223
+ final_width = int(math.ceil(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
224
+ final_height = int(math.ceil(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
225
+
226
+ else:
227
+ final_width = int(round(width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
228
+ final_height = int(round(height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
229
+
230
+ return image.resize((final_width, final_height), Image.Resampling.LANCZOS)
231
+
232
+ def resize_and_crop_to_match(target_image, reference_image):
233
+ """Resizes the target image to match the reference image's dimensions."""
234
+ ref_width, ref_height = reference_image.size
235
+ return target_image.resize((ref_width, ref_height), Image.Resampling.LANCZOS)
236
+
237
+ def crop_to_match(target_image, reference_image):
238
+ """Resizes and center-crops the target image to match the reference image's dimensions."""
239
+ ref_width, ref_height = reference_image.size
240
+ target_width, target_height = target_image.size
241
+ scale = max(ref_width / target_width, ref_height / target_height)
242
+ new_width, new_height = int(target_width * scale), int(target_height * scale)
243
+ resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
244
+ left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
245
+ return resized.crop((left, top, left + ref_width, top + ref_height))
246
+
247
+ def init_view():
248
+ return gr.update(interactive = True)
249
+
250
+ def output_video_change(output_video):
251
+ print('Log output: ' + str(output_video))
252
+ return [gr.update(visible = True)] * 2
253
+
254
+ def generate_video(
255
+ start_image_pil,
256
+ end_image_pil,
257
+ prompt,
258
+ negative_prompt=default_negative_prompt,
259
+ resolution=500000,
260
+ duration_seconds=2.1,
261
+ steps=8,
262
+ guidance_scale=1,
263
+ guidance_scale_2=1,
264
+ seed=42,
265
+ randomize_seed=True,
266
+ progress=gr.Progress(track_tqdm=True)
267
+ ):
268
+ start = time.time()
269
+ allocation_time = 120
270
+ factor = 1
271
+
272
+ if input_image_debug_value[0] is not None or end_image_debug_value[0] is not None or prompt_debug_value[0] is not None or total_second_length_debug_value[0] is not None or allocation_time_debug_value[0] is not None or resolution_debug_value[0] is not None or factor_debug_value[0] is not None:
273
+ start_image_pil = input_image_debug_value[0]
274
+ end_image_pil = end_image_debug_value[0]
275
+ prompt = prompt_debug_value[0]
276
+ duration_seconds = total_second_length_debug_value[0]
277
+ resolution = resolution_debug_value[0]
278
+ factor = factor_debug_value[0]
279
+ allocation_time = allocation_time_debug_value[0]
280
+
281
+ if start_image_pil is None or end_image_pil is None:
282
+ raise gr.Error("Please upload both a start and an end image.")
283
+
284
+ # Step 1: Process the start image to get our target dimensions based on the new rules.
285
+ processed_start_image = process_image_for_video(start_image_pil, resolution)
286
+
287
+ # Step 2: Make the end image match the *exact* dimensions of the processed start image.
288
+ processed_end_image = resize_and_crop_to_match(end_image_pil, processed_start_image)
289
+
290
+ target_height, target_width = processed_start_image.height, processed_start_image.width
291
+
292
+ # Handle seed and frame count
293
+ current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
294
+ num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
295
+
296
+ progress(0.2, desc=f"Generating {num_frames} frames at {target_width}x{target_height} (seed: {current_seed})...")
297
+
298
+ progress(0.1, desc="Preprocessing images...")
299
+ print("Generate a video with the prompt: " + prompt)
300
+ output_frames_list = None
301
+ caught_error = None
302
+ while factor >= 1 and int(allocation_time) > 0:
303
+ try:
304
+ output_frames_list = generate_video_on_gpu(
305
+ start_image_pil,
306
+ end_image_pil,
307
+ prompt,
308
+ negative_prompt,
309
+ int(steps),
310
+ float(guidance_scale),
311
+ float(guidance_scale_2),
312
+ progress,
313
+ allocation_time,
314
+ target_height,
315
+ target_width,
316
+ current_seed,
317
+ (int(((num_frames * factor) - 1) / 4) * 4) + 1,
318
+ processed_start_image,
319
+ processed_end_image
320
+ )
321
+ factor = 0
322
+ caught_error = None
323
+ except BaseException as err:
324
+ print("An exception occurred: " + str(err))
325
+ caught_error = err
326
+ try:
327
+ print('e.message: ' + err.message) # No GPU is currently available for you after 60s
328
+ except Exception as e2:
329
+ print('Failure')
330
+ if not str(err).startswith("No GPU is currently available for you after 60s"):
331
+ factor -= .003
332
+ allocation_time = int(allocation_time) - 1
333
+ except:
334
+ print("An error occurred")
335
+ caught_error = None
336
+ if not str(e).startswith("No GPU is currently available for you after 60s"):
337
+ factor -= .003
338
+ allocation_time = int(allocation_time) - 1
339
+
340
+ if caught_error is not None:
341
+ return [gr.skip(), gr.skip(), gr.skip(), gr.update(value=_error_to_html(caught_error), visible=True), gr.skip()]
342
+
343
+ input_image_debug_value[0] = end_image_debug_value[0] = prompt_debug_value[0] = total_second_length_debug_value[0] = allocation_time_debug_value[0] = factor_debug_value[0] = None
344
+
345
+ progress(0.9, desc="Encoding and saving video...")
346
+
347
+ video_path = 'wan_' + datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.mp4'
348
+
349
+ export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
350
+ set_mp4_comments_imageio_ffmpeg(video_path, f"Prompt: {prompt} | Negative Prompt: {negative_prompt}");
351
+ print("Video exported: " + video_path)
352
+
353
+ progress(1.0, desc="Done!")
354
+ end = time.time()
355
+ secondes = int(end - start)
356
+ minutes = math.floor(secondes / 60)
357
+ secondes = secondes - (minutes * 60)
358
+ hours = math.floor(minutes / 60)
359
+ minutes = minutes - (hours * 60)
360
+ information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
361
+ "The video been generated in " + \
362
+ ((str(hours) + " h, ") if hours != 0 else "") + \
363
+ ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
364
+ str(secondes) + " sec (including " + str(allocation_time) + " seconds of GPU). " + \
365
+ "The video has " + str(int(num_frames * factor)) + " frames. " + \
366
+ "The video resolution is " + str(target_width) + \
367
+ " pixels large and " + str(target_height) + \
368
+ " pixels high, so a resolution of " + f'{target_width * target_height:,}' + " pixels." + \
369
+ " Your prompt is saved into the metadata of the video."
370
+ return [video_path, gr.update(value = video_path, visible = True, interactive = True), current_seed, gr.update(value = information, visible = True), gr.update(interactive = False)]
371
+
372
+ def get_duration(
373
+ start_image_pil,
374
+ end_image_pil,
375
+ prompt,
376
+ negative_prompt,
377
+ steps,
378
+ guidance_scale,
379
+ guidance_scale_2,
380
+ progress,
381
+ allocation_time,
382
+ target_height,
383
+ target_width,
384
+ current_seed,
385
+ num_frames,
386
+ processed_start_image,
387
+ processed_end_image
388
+ ):
389
+ return allocation_time
390
+
391
+ @torch.no_grad()
392
+ @spaces.GPU(duration=get_duration)
393
+ def generate_video_on_gpu(
394
+ start_image_pil,
395
+ end_image_pil,
396
+ prompt,
397
+ negative_prompt,
398
+ steps,
399
+ guidance_scale,
400
+ guidance_scale_2,
401
+ progress,
402
+ allocation_time,
403
+ target_height,
404
+ target_width,
405
+ current_seed,
406
+ num_frames,
407
+ processed_start_image,
408
+ processed_end_image
409
+ ):
410
+ """
411
+ Generates a video by interpolating between a start and end image, guided by a text prompt,
412
+ using the diffusers Wan2.2 pipeline.
413
+ """
414
+
415
+ output_frames_list = pipe(
416
+ image=processed_start_image,
417
+ last_image=processed_end_image,
418
+ prompt=prompt,
419
+ negative_prompt=negative_prompt,
420
+ height=target_height,
421
+ width=target_width,
422
+ num_frames=num_frames,
423
+ guidance_scale=guidance_scale,
424
+ guidance_scale_2=guidance_scale_2,
425
+ num_inference_steps=steps,
426
+ generator=torch.Generator(device="cuda").manual_seed(current_seed),
427
+ ).frames[0]
428
+
429
+ return output_frames_list
430
+
431
+ def export_compiled_transformers_to_zip() -> str:
432
+ """
433
+ Bundle compiled_transformer_1 and compiled_transformer_2 into a zip file and return the file path.
434
+ """
435
+ ct1 = getattr(optimization, "COMPILED_TRANSFORMER_1", None)
436
+ ct2 = getattr(optimization, "COMPILED_TRANSFORMER_2", None)
437
+
438
+ if ct1 is None or ct2 is None:
439
+ raise gr.Error("Compiled transformers are not available yet (compilation may have failed).")
440
+
441
+ payload_1 = ct1.to_serializable_dict()
442
+ payload_2 = ct2.to_serializable_dict()
443
+
444
+ tmp_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
445
+ tmp_zip.close()
446
+
447
+ with zipfile.ZipFile(tmp_zip.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
448
+ # store with torch.save so users can load easily with torch.load()
449
+ buf1 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
450
+ buf1.close()
451
+ torch.save(payload_1, buf1.name)
452
+
453
+ buf2 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
454
+ buf2.close()
455
+ torch.save(payload_2, buf2.name)
456
+
457
+ zf.write(buf1.name, arcname="compiled_transformer_1.pt")
458
+ zf.write(buf2.name, arcname="compiled_transformer_2.pt")
459
+
460
+ # cleanup intermediate .pt
461
+ try:
462
+ os.remove(buf1.name)
463
+ os.remove(buf2.name)
464
+ except:
465
+ pass
466
+
467
+ return tmp_zip.name
468
+
469
+
470
+ # --- 3. Gradio User Interface ---
471
+
472
+
473
+
474
+ js = """
475
+ function createGradioAnimation() {
476
+ window.addEventListener("beforeunload", function(e) {
477
+ if (document.getElementById('dummy_button_id') && !document.getElementById('dummy_button_id').disabled) {
478
+ var confirmationMessage = 'A process is still running. '
479
+ + 'If you leave before saving, your changes will be lost.';
480
+
481
+ (e || window.event).returnValue = confirmationMessage;
482
+ }
483
+ return confirmationMessage;
484
+ });
485
+ return 'Animation created';
486
+ }
487
+ """
488
+
489
+ # Gradio interface
490
+ with gr.Blocks(js=js) as app:
491
+ gr.Markdown("# Wan 2.2 First/Last Frame Video Fast")
492
+ gr.Markdown("Based on the [Wan 2.2 First/Last Frame workflow](https://www.reddit.com/r/StableDiffusion/comments/1me4306/psa_wan_22_does_first_frame_last_frame_out_of_the/), applied to 🧨 Diffusers + [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning) 8-step LoRA")
493
+
494
+ with gr.Row(elem_id="general_items"):
495
+ with gr.Column():
496
+ with gr.Group(elem_id="group_all"):
497
+ with gr.Row():
498
+ start_image = gr.Image(type="pil", label="Start Frame", sources=["upload", "clipboard"])
499
+ # Capture the Tabs component in a variable and assign IDs to tabs
500
+ with gr.Tabs(elem_id="group_tabs") as tabs:
501
+ with gr.TabItem("Upload", id="upload_tab"):
502
+ end_image = gr.Image(type="pil", label="End Frame", sources=["upload", "clipboard"])
503
+ with gr.TabItem("Generate", id="generate_tab"):
504
+ generate_5seconds = gr.Button("Generate scene 5 seconds in the future", elem_id="fivesec")
505
+ gr.Markdown("Generate a custom end-frame with an edit model like [Nano Banana](https://huggingface.co/spaces/multimodalart/nano-banana) or [Qwen Image Edit](https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast)", elem_id="or_item")
506
+ prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images", placeholder="The creature starts to move")
507
+
508
+ with gr.Accordion("Advanced Settings", open=False):
509
+ duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=2.1, label="Video Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
510
+ negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
511
+ resolution = gr.Dropdown([
512
+ ["400,000 px (working)", 400000],
513
+ ["465,920 px (working)", 465920],
514
+ ["495,616 px (working)", 495616],
515
+ ["500,000 px (working)", 500000],
516
+ ["600,000 px (working)", 600000],
517
+ ["700,000 px (working)", 700000],
518
+ ["800,000 px (working)", 800000],
519
+ ["900,000 px (working)", 900000],
520
+ ["1,000,000 px (working)", 1000000],
521
+ ["1,100,000 px (untested)", 1100000],
522
+ ["1,200,000 px (untested)", 1200000],
523
+ ["1,300,000 px (untested)", 1300000],
524
+ ["1,400,000 px (untested)", 1400000],
525
+ ["1,500,000 px (untested)", 1500000]
526
+ ], value=465920, label="Resolution (width x height)", info="Less if the image is smaller")
527
+ steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
528
+ guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - high noise")
529
+ guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - low noise")
530
+ with gr.Row():
531
+ seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
532
+ randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
533
+
534
+ generate_button = gr.Button("🚀 Generate Video", variant="primary")
535
+ dummy_button = gr.Button(elem_id = "dummy_button_id", visible = False, interactive = False)
536
+
537
+ with gr.Column():
538
+ output_video = gr.Video(elem_id="output_id", label="Generated Video", autoplay = True, loop = True)
539
+ download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
540
+ video_information = gr.HTML(value = "")
541
+
542
+ with gr.Accordion("🔧 Compilation artifacts (advanced)", open=False):
543
+ gr.Markdown(
544
+ "Télécharge les artefacts compilés AOTInductor générés au démarrage (transformer + transformer_2)."
545
+ )
546
+ export_btn = gr.Button("📦 Préparer l'archive des transformers compilés")
547
+ compiled_download = gr.DownloadButton(label="⬇️ Télécharger compiled_transformers.zip", interactive=False)
548
+
549
+ def _build_and_enable_download():
550
+ path = export_compiled_transformers_to_zip()
551
+ return gr.update(value=path, interactive=True)
552
+
553
+ export_btn.click(fn=_build_and_enable_download, inputs=None, outputs=compiled_download)
554
+
555
+ # Main video generation button
556
+ ui_inputs = [
557
+ start_image,
558
+ end_image,
559
+ prompt,
560
+ negative_prompt_input,
561
+ resolution,
562
+ duration_seconds_input,
563
+ steps_slider,
564
+ guidance_scale_input,
565
+ guidance_scale_2_input,
566
+ seed_input,
567
+ randomize_seed_checkbox
568
+ ]
569
+ ui_outputs = [output_video, download_button, seed_input, video_information, dummy_button]
570
+
571
+ generate_button.click(fn = init_view, inputs = [], outputs = [dummy_button], queue = False, show_progress = False).success(
572
+ fn = generate_video,
573
+ inputs = ui_inputs,
574
+ outputs = ui_outputs
575
+ )
576
+
577
+ generate_5seconds.click(
578
+ fn=switch_to_upload_tab,
579
+ inputs=None,
580
+ outputs=[tabs]
581
+ ).then(
582
+ fn=lambda img: generate_end_frame(img, "this image is a still frame from a movie. generate a new frame with what happens on this scene 5 seconds in the future"),
583
+ inputs=[start_image],
584
+ outputs=[end_image]
585
+ ).success(
586
+ fn=generate_video,
587
+ inputs=ui_inputs,
588
+ outputs=ui_outputs
589
+ )
590
+
591
+ output_video.change(
592
+ fn=output_video_change,
593
+ inputs=[output_video],
594
+ outputs=[download_button, video_information],
595
+ js="document.getElementById('download_btn').click()"
596
+ ).then(
597
+ fn=lambda: None,
598
+ inputs=[],
599
+ outputs=[],
600
+ js="document.getElementById('output_id')?.scrollIntoView({behavior: 'smooth', block: 'start', inline: 'nearest'})"
601
+ )
602
+
603
+ with gr.Row(visible=False):
604
+ gr.Examples(
605
+ examples=[["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."]],
606
+ inputs=[start_image, end_image, prompt],
607
+ outputs=ui_outputs,
608
+ fn=generate_video,
609
+ run_on_click=True,
610
+ cache_examples=True,
611
+ cache_mode='lazy',
612
+ )
613
+ prompt_debug=gr.Textbox(label="Prompt Debug")
614
+ input_image_debug=gr.Image(type="pil", label="Image Debug")
615
+ end_image_debug=gr.Image(type="pil", label="End Image Debug")
616
+ total_second_length_debug=gr.Slider(label="Duration Debug", minimum=1, maximum=120, value=5, step=0.1)
617
+ resolution_debug = gr.Dropdown([
618
+ ["400,000 px", 400000],
619
+ ["465,920 px", 465920],
620
+ ["495,616 px", 495616],
621
+ ["500,000 px", 500000],
622
+ ["600,000 px", 600000],
623
+ ["700,000 px", 700000],
624
+ ["800,000 px", 800000],
625
+ ["900,000 px", 900000],
626
+ ["1,000,000 px", 1000000],
627
+ ["1,100,000 px", 1100000],
628
+ ["1,200,000 px", 1200000],
629
+ ["1,300,000 px", 1300000],
630
+ ["1,400,000 px", 1400000],
631
+ ["1,500,000 px", 1500000]
632
+ ], value=500000, label="Resolution Debug")
633
+ factor_debug=gr.Slider(label="Factor Debug", minimum=1, maximum=100, value=3.2, step=0.1)
634
+ allocation_time_debug=gr.Slider(label="Allocation Debug", minimum=1, maximum=60 * 20, value=720, step=1)
635
+
636
+ def handle_field_debug_change(
637
+ input_image_debug_data,
638
+ end_image_debug_data,
639
+ prompt_debug_data,
640
+ total_second_length_debug_data,
641
+ resolution_debug_data,
642
+ factor_debug_data,
643
+ allocation_time_debug_data
644
+ ):
645
+ input_image_debug_value[0] = input_image_debug_data
646
+ end_image_debug_value[0] = end_image_debug_data
647
+ prompt_debug_value[0] = prompt_debug_data
648
+ total_second_length_debug_value[0] = total_second_length_debug_data
649
+ resolution_debug_value[0] = resolution_debug_data
650
+ factor_debug_value[0] = factor_debug_data
651
+ allocation_time_debug_value[0] = allocation_time_debug_data
652
+ return []
653
+
654
+ inputs_debug=[input_image_debug, end_image_debug, prompt_debug, total_second_length_debug, resolution_debug, factor_debug, allocation_time_debug]
655
+
656
+ input_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
657
+ end_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
658
+ prompt_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
659
+ total_second_length_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
660
+ resolution_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
661
+ factor_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
662
+ allocation_time_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
663
+
664
+ gr.Examples(
665
+ label = "Examples from demo",
666
+ examples = [
667
+ ["poli_tower.png", "tower_takes_off.png", "The man turns around."],
668
+ ["ugly_sonic.jpeg", "squatting_sonic.png", "पात्रं क्षेपणास्त्रं चकमाति।"],
669
+ ["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."],
670
+ ],
671
+ inputs = [start_image, end_image, prompt],
672
+ outputs = ui_outputs,
673
+ fn = generate_video,
674
+ cache_examples = False,
675
+ )
676
+
677
+ if __name__ == "__main__":
678
  app.launch(mcp_server=True, share=True)