Update app_fast.py

app_fast.py

@@ -1,216 +1,247 @@
-# PyTorch 2.8 (temporary hack)
-import os
-os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
-#os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu128 "torch<2.9" spaces')
-
-# Actual demo code
 import spaces
 import torch
-from diffusers import WanPipeline, AutoencoderKLWan
-from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
-from diffusers.utils.export_utils import export_to_video
+from diffusers import AutoencoderKLWan, WanPipeline, WanImageToVideoPipeline, UniPCMultistepScheduler
+from diffusers.utils import export_to_video
 import gradio as gr
 import tempfile
+from huggingface_hub import snapshot_download, hf_hub_download
 import numpy as np
 from PIL import Image
 import random
-import gc
-from optimization import optimize_pipeline_
-from accelerate import Accelerator
-
-
-def get_memory_gb(device):
-    """Get current allocated memory in GB"""
-    return torch.cuda.memory_allocated(device) / 1024**3
-
-accelerator = Accelerator()
-device = accelerator.device
-# Set up device and dtype using accelerator
-
-initial_memory = get_memory_gb(device)
-print(f"Initial memory: {initial_memory:.2f}GB")
+import os
 
+# ---------------------------
+# Configuration - edit these
+# ---------------------------
+# This should be the HF repo id that contains the model files (configs + safetensors)
+# e.g. "username/wan2.2-t2v-rapid-aio-v6" or the repo where v6/wan2.2... is stored.
+MODEL_REPO = "Phr00t/WAN2.2-14B-Rapid-AllInOne"
 
+# If your model files are nested under a folder in the repo (e.g. "v6/wan2.2-t2v-rapid-aio-v6.safetensors"),
+# set MODEL_FILE to that relative path (optional). If the repo is already standard diffusers layout, leave None.
+MODEL_FILE = "v6/wan2.2-t2v-rapid-aio-v6.safetensors"  # or None
 
-MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+# Preferred dtype for model weights on GPU (float16 or bfloat16). If your GPU doesn't support bfloat16 use float16.
+PREFERRED_DTYPE = torch.bfloat16
 
-LANDSCAPE_WIDTH = 832
-LANDSCAPE_HEIGHT = 480
+# ---------------------------
+# Helper: prepare local model
+# ---------------------------
+def prepare_model_local(repo_id: str, model_file: str | None = None):
+    """
+    Downloads the repo snapshot and returns local_model_dir.
+    If model_file is provided we ensure it's present (hf_hub_download fallback).
+    """
+    print(f"[prepare_model_local] snapshot_download(repo_id={repo_id})")
+    local_dir = snapshot_download(repo_id)
+    # if model_file is specified, ensure it's present locally (some repos may not be in root)
+    if model_file:
+        candidate = os.path.join(local_dir, model_file)
+        if not os.path.exists(candidate):
+            # try to download the single file into the local_dir
+            try:
+                print(f"[prepare_model_local] hf_hub_download(repo_id={repo_id}, filename={model_file})")
+                path = hf_hub_download(repo_id=repo_id, filename=model_file, cache_dir=None)
+                # Move / copy into the snapshot folder for consistent load paths (optional)
+                # but usually snapshot_download should have returned the file already
+                print(f"[prepare_model_local] got file at {path}")
+            except Exception as e:
+                print(f"[prepare_model_local] WARNING: couldn't find {model_file} in repo: {e}")
+        else:
+            print(f"[prepare_model_local] found {candidate}")
+    return local_dir
+
+# ---------------------------
+# Download & load model
+# ---------------------------
+local_model_dir = prepare_model_local(MODEL_REPO, MODEL_FILE)
+
+# Try to load a VAE if present in subfolder "vae"
+vae = None
+vae_dir = os.path.join(local_model_dir, "vae")
+try:
+    if os.path.isdir(vae_dir):
+        print("[main] Loading VAE from", vae_dir)
+        vae = AutoencoderKLWan.from_pretrained(vae_dir, torch_dtype=torch.float32)
+    else:
+        # If repo layout packs everything into root, WanPipeline.from_pretrained will handle it
+        print("[main] No vae/ subfolder detected - will rely on pipeline repo layout.")
+except Exception as e:
+    print("[main] VAE load failed, continuing without explicit VAE:", e)
+    vae = None
+
+# Attempt model load with preferred dtype; fall back to float32 if necessary
+def load_pipeline_with_fallback(pipeline_cls, repo_or_dir, vae_obj=None, dtype=PREFERRED_DTYPE):
+    try:
+        print(f"[load_pipeline_with_fallback] trying dtype={dtype}")
+        return pipeline_cls.from_pretrained(repo_or_dir, vae=vae_obj, torch_dtype=dtype)
+    except Exception as e:
+        print(f"[load_pipeline_with_fallback] failed with dtype={dtype}: {e}")
+        print("[load_pipeline_with_fallback] retrying with float32")
+        return pipeline_cls.from_pretrained(repo_or_dir, vae=vae_obj, torch_dtype=torch.float32)
+
+# Load pipelines
+text_to_video_pipe = load_pipeline_with_fallback(WanPipeline, local_model_dir, vae_obj=vae)
+image_to_video_pipe = load_pipeline_with_fallback(WanImageToVideoPipeline, local_model_dir, vae_obj=vae)
+
+# Adjust scheduler and move to GPU
+for pipe in [text_to_video_pipe, image_to_video_pipe]:
+    try:
+        # keep your custom scheduler tweak
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
+    except Exception as e:
+        print("[main] Warning: couldn't replace scheduler:", e)
+    # move to CUDA
+    pipe.to("cuda")
+    # optionally enable attention slicing / memory optimizations if needed:
+    try:
+        pipe.enable_attention_slicing()
+    except Exception:
+        pass
+
+# -------------------------------------------------------------------------
+# (rest of your code remains the same)
+# -------------------------------------------------------------------------
+# Constants
+MOD_VALUE = 32
+DEFAULT_H_SLIDER_VALUE = 896
+DEFAULT_W_SLIDER_VALUE = 896
+NEW_FORMULA_MAX_AREA = 720 * 1024
+SLIDER_MIN_H, SLIDER_MAX_H = 256, 1024
+SLIDER_MIN_W, SLIDER_MAX_W = 256, 1024
 MAX_SEED = np.iinfo(np.int32).max
-
-FIXED_FPS = 16
-MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 81
-
-MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
-MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
-
-vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.2-T2V-A14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
-pipe = WanPipeline.from_pretrained(MODEL_ID,
-    transformer=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
-        subfolder='transformer',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    transformer_2=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
-        subfolder='transformer_2',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    vae=vae,
-    torch_dtype=torch.bfloat16,
-).to('cuda')
-
-
-for i in range(3): 
-    gc.collect()
-    torch.cuda.synchronize() 
-    torch.cuda.empty_cache()
-
-optimize_pipeline_(pipe,
-    prompt='prompt',
-    height=LANDSCAPE_HEIGHT,
-    width=LANDSCAPE_WIDTH,
-    num_frames=MAX_FRAMES_MODEL,
-)
-
-
-default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
-default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
-
-
-def get_duration(
-    prompt,
-    negative_prompt,
-    duration_seconds,
-    guidance_scale,
-    guidance_scale_2,
-    steps,
-    seed,
-    randomize_seed,
-    progress,
-):
-    return steps * 15
+FIXED_FPS = 24
+MIN_FRAMES_MODEL = 25
+MAX_FRAMES_MODEL = 193
+
+default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
+default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature"
+
+def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area, min_slider_h, max_slider_h, min_slider_w, max_slider_w, default_h, default_w):
+    orig_w, orig_h = pil_image.size
+    if orig_w <= 0 or orig_h <= 0:
+        return default_h, default_w
+    aspect_ratio = orig_h / orig_w
+
+    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
+    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
+    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
+    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
+
+    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
+    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
+
+    return new_h, new_w
+
+def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
+    if uploaded_pil_image is None:
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+    try:
+        new_h, new_w = _calculate_new_dimensions_wan(
+            uploaded_pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
+            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
+            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
+        )
+        return gr.update(value=new_h), gr.update(value=new_w)
+    except Exception as e:
+        gr.Warning("Error attempting to calculate new dimensions")
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+        
+def get_duration(input_image, prompt, height, width, 
+                   negative_prompt, duration_seconds,
+                   guidance_scale, steps,
+                   seed, randomize_seed, 
+                   progress):
+    if steps > 4 and duration_seconds > 4:
+        return 90
+    elif steps > 4 or duration_seconds > 4:
+        return 75
+    else:
+        return 60
 
 @spaces.GPU(duration=get_duration)
-def generate_video(
-    prompt,
-    negative_prompt=default_negative_prompt,
-    duration_seconds = MAX_DURATION,
-    guidance_scale = 1,
-    guidance_scale_2 = 3,
-    steps = 4,
-    seed = 42,
-    randomize_seed = False,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generate a video from a text prompt using the Wan 2.2 14B T2V model with Lightning LoRA.
-    
-    This function takes an input prompt and generates a video animation based on the provided
-    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Text-to-Video model with Lightning LoRA
-    for fast generation in 4-8 steps.
-    
-    Args:
-        prompt (str): Text prompt describing the desired animation or motion.
-        negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
-            Defaults to default_negative_prompt (contains unwanted visual artifacts).
-        duration_seconds (float, optional): Duration of the generated video in seconds.
-            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
-        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
-            Defaults to 4. Range: 1-30.
-        seed (int, optional): Random seed for reproducible results. Defaults to 42.
-            Range: 0 to MAX_SEED (2147483647).
-        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
-            Defaults to False.
-        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
-    
-    Returns:
-        tuple: A tuple containing:
-            - video_path (str): Path to the generated video file (.mp4)
-            - current_seed (int): The seed used for generation (useful when randomize_seed=True)
-    
-    Raises:
-        gr.Error: If input_image is None (no image uploaded).
-    
-    Note:
-        - The function automatically resizes the input image to the target dimensions
-        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
-        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
-        - The function uses GPU acceleration via the @spaces.GPU decorator
-        - Generation time varies based on steps and duration (see get_duration function)
-    """
-    
+def generate_video(input_image, prompt, height, width, negative_prompt=default_negative_prompt, duration_seconds=2, guidance_scale=0, steps=4, seed=44, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
+    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
+    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
+
     num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
-    output_frames_list = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=480,
-        width=832,
-        num_frames=num_frames,
-        guidance_scale=float(guidance_scale),
-        guidance_scale_2=float(guidance_scale_2),
-        num_inference_steps=int(steps),
-        generator=torch.Generator(device="cuda").manual_seed(current_seed),
-    ).frames[0]
+    if input_image is not None:
+        resized_image = input_image.resize((target_w, target_h))
+        with torch.inference_mode():
+            output_frames_list = image_to_video_pipe(
+                image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
+                height=target_h, width=target_w, num_frames=num_frames,
+                guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
+                generator=torch.Generator(device="cuda").manual_seed(current_seed)
+            ).frames[0]
+    else:
+        with torch.inference_mode():
+            output_frames_list = text_to_video_pipe(
+                prompt=prompt, negative_prompt=negative_prompt,
+                height=target_h, width=target_w, num_frames=num_frames,
+                guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
+                generator=torch.Generator(device="cuda").manual_seed(current_seed)
+            ).frames[0]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
-
     export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
-
     return video_path, current_seed
 
+# UI (unchanged)
 with gr.Blocks() as demo:
-    gr.Markdown("# Fast 4 steps Wan 2.2 T2V (14B) with Lightning LoRA")
-    gr.Markdown("run Wan 2.2 in just 4-8 steps, with [Wan 2.2 Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
+    gr.Markdown("# Fast Wan 2.2 TI2V 5B Demo")
+    gr.Markdown("""This Demo is using your Wan2.2 model from the specified HF repo.""")
     with gr.Row():
         with gr.Column():
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
-            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
-            
+            input_image_component = gr.Image(type="pil", label="Input Image (optional, auto-resized to target H/W)")
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
+            duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
+
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                 seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                 randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps") 
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
-                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3, label="Guidance Scale 2 - low noise stage")
-
+                with gr.Row():
+                    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})")
+                    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})")
+                steps_slider = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps")
+                guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.0, label="Guidance Scale")
             generate_button = gr.Button("Generate Video", variant="primary")
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-    
+
+    input_image_component.upload(
+        fn=handle_image_upload_for_dims_wan,
+        inputs=[input_image_component, height_input, width_input],
+        outputs=[height_input, width_input]
+    )
+
+    input_image_component.clear(
+        fn=handle_image_upload_for_dims_wan,
+        inputs=[input_image_component, height_input, width_input],
+        outputs=[height_input, width_input]
+    )
+
     ui_inputs = [
-        prompt_input,
+        input_image_component, prompt_input, height_input, width_input,
         negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, guidance_scale_2_input, steps_slider, seed_input, randomize_seed_checkbox
+        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
     ]
     generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
     gr.Examples(
         examples=[ 
-            [
-                "POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood.",
-            ],
-            [
-                "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
-            ],
-            [
-                "A cinematic shot of a boat sailing on a calm sea at sunset.",
-            ],
-            [
-                "Drone footage flying over a futuristic city with flying cars.",
-            ],
+            [None, "A person eating spaghetti", 1024, 720],
+            ["cat.png", "The cat removes the glasses from its eyes.", 1088, 800],
+            [None, "a penguin playfully dancing in the snow, Antarctica", 1024, 720],
+            ["peng.png", "a penguin running towards camera joyfully, Antarctica", 896, 512],
         ],
-        inputs=[prompt_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
+        inputs=[input_image_component, prompt_input, height_input, width_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(mcp_server=True)
+    demo.queue().launch()