Update app.py

app.py

@@ -1,223 +1,347 @@
-# CACHE_BUSTER = "2025-12-15-ROBUST"
-# ======================================================
-# Wan I2V – ROBUST, CLEAN, HF-SPACES SAFE IMPLEMENTATION
-# ======================================================
-# Goals:
-# - Zero syntax errors (Python 3 only)
-# - No Gradio hot-reload issues
-# - Robust video export (diffusers → imageio → ffmpeg → opencv)
-# - Safe long-video chunking
-# - Clear logging
-# - HF Spaces compatible (ZeroGPU, CPU, CUDA)
-
 import os
-import uuid
-import time
-import shutil
-import logging
-from typing import List, Tuple
-
+import spaces
+import torch
+from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
+from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
+from diffusers.utils.export_utils import export_to_video
+import gradio as gr
+import tempfile
 import numpy as np
 from PIL import Image
-import gradio as gr
-
-# ------------------------------------------------------
-# Optional backends
-# ------------------------------------------------------
-try:
-    from diffusers.utils import export_to_video as diffusers_export_to_video
-    DIFFUSERS_AVAILABLE = True
-except Exception:
-    DIFFUSERS_AVAILABLE = False
-
-try:
-    import imageio
-    IMAGEIO_AVAILABLE = True
-except Exception:
-    IMAGEIO_AVAILABLE = False
-
-try:
-    import cv2
-    CV2_AVAILABLE = True
-except Exception:
-    CV2_AVAILABLE = False
-
-# ------------------------------------------------------
-# Logging
-# ------------------------------------------------------
-logging.basicConfig(
-    level=logging.INFO,
-    format="[%(asctime)s] [%(levelname)s] %(message)s",
+import random
+import gc
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
+import aoti
+
+# =========================================================
+# MODEL CONFIGURATION
+# =========================================================
+MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+MAX_DIM = 832
+MIN_DIM = 480
+SQUARE_DIM = 640
+MULTIPLE_OF = 16
+MAX_SEED = np.iinfo(np.int32).max
+
+FIXED_FPS = 16
+# We will generate in chunks of ~5 seconds (81 frames) to reach 20s
+CHUNK_DURATION = 5.0 
+TOTAL_DURATION_TARGET = 20.0
+
+# =========================================================
+# LOAD PIPELINE
+# =========================================================
+print("Loading pipeline components...")
+
+# Load models in bfloat16
+transformer = WanTransformer3DModel.from_pretrained(
+    MODEL_ID,
+    subfolder="transformer",
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN
 )
-logger = logging.getLogger("wan_i2v")
-
-# ------------------------------------------------------
-# Frame normalization
-# ------------------------------------------------------
-
-def normalize_frame(frame: np.ndarray) -> np.ndarray:
-    frame = np.asarray(frame)
-
-    if frame.ndim == 4 and frame.shape[0] == 1:
-        frame = frame[0]
-
-    if frame.ndim == 2:
-        frame = np.stack([frame] * 3, axis=-1)
-
-    if frame.ndim == 3 and frame.shape[0] in (1, 3):
-        frame = np.transpose(frame, (1, 2, 0))
-
-    if frame.dtype != np.uint8:
-        if frame.min() < 0:
-            frame = (frame + 1.0) / 2.0
-        frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
-
-    return frame
-
-# ------------------------------------------------------
-# Robust video export
-# ------------------------------------------------------
-
-def export_video(frames: List[np.ndarray], out_path: str, fps: int) -> str:
-    frames = [normalize_frame(f) for f in frames]
-
-    # 1. diffusers
-    if DIFFUSERS_AVAILABLE:
-        try:
-            diffusers_export_to_video(frames, out_path, fps=fps)
-            return out_path
-        except Exception as e:
-            logger.warning(f"Diffusers export failed: {e}")
-
-    # 2. imageio
-    if IMAGEIO_AVAILABLE:
-        try:
-            writer = imageio.get_writer(out_path, fps=fps, macro_block_size=None)
-            for f in frames:
-                writer.append_data(f)
-            writer.close()
-            return out_path
-        except Exception as e:
-            logger.warning(f"ImageIO export failed: {e}")
-
-    # 3. ffmpeg raw pipe
-    ffmpeg = shutil.which("ffmpeg")
-    if ffmpeg:
-        try:
-            import subprocess
-            h, w = frames[0].shape[:2]
-            cmd = [
-                ffmpeg, "-y",
-                "-f", "rawvideo",
-                "-pix_fmt", "rgb24",
-                "-s", f"{w}x{h}",
-                "-r", str(fps),
-                "-i", "-",
-                "-an",
-                "-c:v", "libx264",
-                "-pix_fmt", "yuv420p",
-                out_path,
-            ]
-            p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
-            for f in frames:
-                p.stdin.write(f.tobytes())
-            p.stdin.close()
-            p.wait()
-            return out_path
-        except Exception as e:
-            logger.warning(f"FFmpeg export failed: {e}")
-
-    # 4. OpenCV
-    if CV2_AVAILABLE:
-        try:
-            h, w = frames[0].shape[:2]
-            writer = cv2.VideoWriter(
-                out_path,
-                cv2.VideoWriter_fourcc(*"mp4v"),
-                float(fps),
-                (w, h),
-            )
-            for f in frames:
-                writer.write(cv2.cvtColor(f, cv2.COLOR_RGB2BGR))
-            writer.release()
-            return out_path
-        except Exception as e:
-            logger.warning(f"OpenCV export failed: {e}")
-
-    raise RuntimeError("All video export backends failed")
-
-# ------------------------------------------------------
-# Dummy inference (SAFE placeholder)
-# ------------------------------------------------------
-
-def infer_frames(image: Image.Image, num_frames: int) -> List[np.ndarray]:
-    base = np.asarray(image.convert("RGB"))
-    frames = []
-    for i in range(num_frames):
-        f = base.copy()
-        f[:, :, 0] = np.clip(f[:, :, 0] + i, 0, 255)
-        frames.append(f)
-    return frames
-
-# ------------------------------------------------------
-# Long-video generator
-# ------------------------------------------------------
-
-def generate_video(
-    image: Image.Image,
-    total_frames: int,
-    fps: int,
-    seed: int,
-) -> Tuple[str, int]:
-
-    if image is None:
-        raise gr.Error("Please upload an input image")
-
-    np.random.seed(int(seed))
 
-    CHUNK = 80
-    frames_all: List[np.ndarray] = []
-    remaining = total_frames
+transformer_2 = WanTransformer3DModel.from_pretrained(
+    MODEL_ID,
+    subfolder="transformer_2",
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN
+)
 
-    while remaining > 0:
-        take = min(CHUNK, remaining)
-        frames_all.extend(infer_frames(image, take))
-        remaining -= take
+print("Assembling pipeline...")
+pipe = WanImageToVideoPipeline.from_pretrained(
+    MODEL_ID,
+    transformer=transformer,
+    transformer_2=transformer_2,
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN
+)
 
-    out_path = f"/tmp/wan_{uuid.uuid4().hex}.mp4"
-    export_video(frames_all, out_path, fps)
+print("Moving to CUDA...")
+pipe = pipe.to("cuda")
 
-    return out_path, seed
+# =========================================================
+# LOAD LORA ADAPTERS
+# =========================================================
+print("Loading LoRA adapters...")
+try:
+    pipe.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+        adapter_name="lightx2v"
+    )
+    pipe.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+        adapter_name="lightx2v_2",
+        load_into_transformer_2=True
+    )
 
-# ------------------------------------------------------
-# Gradio UI (single demo object – avoids hot-reload bug)
-# ------------------------------------------------------
+    pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
+    pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
+    pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
+    pipe.unload_lora_weights()
+    print("LoRA loaded and fused successfully.")
+except Exception as e:
+    print(f"Warning: Failed to load LoRA. Continuing without it. Error: {e}")
+
+# =========================================================
+# QUANTIZATION & AOT OPTIMIZATION
+# =========================================================
+print("Applying quantization...")
+torch.cuda.empty_cache()
+gc.collect()
 
-demo = gr.Blocks(title="Wan I2V Robust")
+try:
+    quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+    quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+    quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+
+    print("Loading AOTI blocks...")
+    aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
+    aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
+except Exception as e:
+    print(f"Warning: Quantization/AOTI failed. Running in standard mode might OOM. Error: {e}")
+
+# =========================================================
+# DEFAULT PROMPTS
+# =========================================================
+default_prompt_i2v = "Make this image come alive with dynamic, cinematic human motion. Create smooth, natural, lifelike animation with fluid transitions, expressive body movement, realistic physics, and elegant camera flow. Deliver a polished, high-quality motion style that feels immersive, artistic, and visually captivating."
+
+default_negative_prompt = (
+    "low quality, worst quality, motion artifacts, unstable motion, jitter, frame jitter, wobbling limbs, motion distortion, inconsistent movement, robotic movement, animation-like motion, awkward transitions, incorrect body mechanics, unnatural posing, off-balance poses, broken motion paths, frozen frames, duplicated frames, frame skipping, warped motion, stretching artifacts bad anatomy, incorrect proportions, deformed body, twisted torso, broken joints, dislocated limbs, distorted neck, unnatural spine curvature, malformed hands, extra fingers, missing fingers, fused fingers, distorted legs, extra limbs, collapsed feet, floating feet, foot sliding, foot jitter, backward walking, unnatural gait blurry details, long exposure blur, ghosting, shadow trails, smearing, washed-out colors, overexposure, underexposure, excessive contrast, blown highlights, poorly rendered clothing, fabric glitches, texture warping, clothing merging with body, incorrect cloth physics ugly background, cluttered scene, crowded background, random objects, unwanted text, subtitles, logos, graffiti, grain, noise, static artifacts, compression noise, jpeg artifacts, image-like stillness, painting-like look, cartoon texture, low-resolution textures"
+)
 
-with demo:
-    gr.Markdown("## Wan I2V – Robust Stable Build")
+# =========================================================
+# IMAGE RESIZING LOGIC
+# =========================================================
+def resize_image(image: Image.Image) -> Image.Image:
+    width, height = image.size
+    if width == height:
+        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)
+    
+    aspect_ratio = width / height
+    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
+    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM
+
+    image_to_resize = image
+    if aspect_ratio > MAX_ASPECT_RATIO:
+        crop_width = int(round(height * MAX_ASPECT_RATIO))
+        left = (width - crop_width) // 2
+        image_to_resize = image.crop((left, 0, left + crop_width, height))
+    elif aspect_ratio < MIN_ASPECT_RATIO:
+        crop_height = int(round(width / MIN_ASPECT_RATIO))
+        top = (height - crop_height) // 2
+        image_to_resize = image.crop((0, top, width, top + crop_height))
+    
+    if width > height:
+        target_w = MAX_DIM
+        target_h = int(round(target_w / aspect_ratio))
+    else:
+        target_h = MAX_DIM
+        target_w = int(round(target_h * aspect_ratio))
+        
+    final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF
+    final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF
+    
+    final_w = max(MIN_DIM, min(MAX_DIM, final_w))
+    final_h = max(MIN_DIM, min(MAX_DIM, final_h))
+    
+    return image_to_resize.resize((final_w, final_h), Image.LANCZOS)
+
+def get_num_frames(duration_seconds: float):
+    return 1 + int(np.clip(int(round(duration_seconds * FIXED_FPS)), 8, 300))
+
+# =========================================================
+# MAIN GENERATION FUNCTION (REWRITTEN FOR LONG VIDEO)
+# =========================================================
+@spaces.GPU(duration=300) # Increased timeout for long generation
+def generate_video(
+    input_image_path, 
+    prompt,
+    steps=4,
+    negative_prompt=default_negative_prompt,
+    duration_seconds=20.0, # Defaulting to 20s
+    guidance_scale=1,
+    guidance_scale_2=1,
+    seed=42,
+    randomize_seed=False,
+    progress=gr.Progress(track_tqdm=True),
+):
+    # Cleanup memory
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    try:
+        # 1. Validation checks
+        if not input_image_path:
+            raise gr.Error("Please upload an input image.")
+        if not os.path.exists(input_image_path):
+             raise gr.Error("Image file not found! Please re-upload the image.")
+
+        # 2. Setup
+        original_input_image = Image.open(input_image_path).convert("RGB")
+        current_input_image = resize_image(original_input_image)
+        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+        
+        # Determine number of iterations needed for 20s
+        # 20s / 5s chunks = 4 iterations
+        chunk_duration = 5.0 
+        total_duration = float(duration_seconds)
+        iterations = int(np.ceil(total_duration / chunk_duration))
+        
+        print(f"Starting Long Video Generation: {total_duration}s ({iterations} iterations)")
+        
+        all_frames = []
+        
+        # 3. The Autoregressive Loop
+        for i in range(iterations):
+            progress(i / iterations, desc=f"Generating Part {i+1} of {iterations}...")
+            
+            # Calculate frames for this chunk
+            num_frames = get_num_frames(chunk_duration)
+            
+            print(f"--- Generative Pass {i+1}: Seed {current_seed} ---")
+            
+            output_frames_list = pipe(
+                image=current_input_image,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=current_input_image.height,
+                width=current_input_image.width,
+                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]
+            
+            # Store frames
+            # If this is not the first chunk, we drop the first frame of the new chunk
+            # because it is (theoretically) identical to the last frame of the previous chunk
+            if i > 0:
+                all_frames.extend(output_frames_list[1:])
+            else:
+                all_frames.extend(output_frames_list)
+            
+            # Prepare for next iteration
+            # The last frame of this video becomes the input for the next video
+            # We convert the numpy/PIL frame back to PIL for the pipeline
+            last_frame = output_frames_list[-1]
+            current_input_image = last_frame
+            
+            # Optional: Slightly shift seed per chunk to prevent looping artifacts, 
+            # or keep it same for consistency. Keeping same is usually safer for style.
+            
+            # Cleanup per chunk
+            del output_frames_list
+            torch.cuda.empty_cache()
+
+        # 4. Save Final Long Video
+        print(f"Stitching {len(all_frames)} frames...")
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+            video_path = tmpfile.name
+        
+        export_to_video(all_frames, video_path, fps=FIXED_FPS)
+        
+        # Final Cleanup
+        del all_frames
+        del current_input_image
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        return video_path, current_seed
+
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        raise gr.Error(f"Generation failed: {str(e)}")
+
+# =========================================================
+# GRADIO UI
+# =========================================================
+
+# Google Analytics Script
+ga_script = """
+<script async src="https://www.googletagmanager.com/gtag/js?id=G-1TD40BVM04"></script>
+<script>
+  window.dataLayer = window.dataLayer || [];
+  function gtag(){dataLayer.push(arguments);}
+  gtag('js', new Date());
+
+  gtag('config', 'G-1TD40BVM04');
+</script>    
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), head=ga_script) as demo:
+    
+    # --- PROFESSIONAL YOUTUBE EMBED SECTION ---
+    gr.HTML("""
+    <div style="background: linear-gradient(135deg, #b90000 0%, #ff0000 100%); color: white; padding: 25px; border-radius: 16px; text-align: center; margin-bottom: 25px; box-shadow: 0 10px 30px rgba(185, 0, 0, 0.3);">
+        <div style="display: flex; align-items: center; justify-content: center; gap: 25px; flex-wrap: wrap; margin-bottom: 20px;">
+            <div style="display: flex; align-items: center; gap: 15px;">
+                <div style="background: white; width: 50px; height: 50px; border-radius: 50%; display: flex; align-items: center; justify-content: center; box-shadow: 0 4px 8px rgba(0,0,0,0.2);">
+                    <span style="font-size: 24px;">▶️</span>
+                </div>
+                <div style="text-align: left;">
+                    <h3 style="margin: 0; font-weight: 800; font-size: 22px; letter-spacing: 0.5px;">Imagination Engineering</h3>
+                    <p style="margin: 4px 0 0 0; opacity: 0.95; font-size: 14px; font-weight: 400;">Mastering AI & Creative Tech</p>
+                </div>
+            </div>
+            <a href="https://www.youtube.com/@ImaginationEngineering" target="_blank" style="text-decoration: none;">
+                <button style="background-color: white; color: #cc0000; border: none; padding: 10px 28px; border-radius: 30px; font-weight: 700; cursor: pointer; transition: transform 0.2s, box-shadow 0.2s; font-size: 15px; box-shadow: 0 4px 12px rgba(0,0,0,0.2);">
+                    SUBSCRIBE & WATCH 📺
+                </button>
+            </a>
+        </div>
+    </div>
+    """)
 
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Input Image")
-        fps = gr.Slider(8, 60, value=24, step=1, label="FPS")
-
-    total_frames = gr.Slider(16, 480, value=80, step=1, label="Total Frames")
-    seed = gr.Number(value=1234, label="Seed")
-
-    run_btn = gr.Button("Generate Video")
-    output_video = gr.Video(label="Output Video")
-    output_seed = gr.Number(label="Seed Used")
-
-    run_btn.click(
-        fn=generate_video,
-        inputs=[input_image, total_frames, fps, seed],
-        outputs=[output_video, output_seed],
-    )
-
-# ------------------------------------------------------
-# Launch
-# ------------------------------------------------------
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="filepath", label="Input Image", elem_id="input_image")
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
+            negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                # Set default to 20 seconds
+                duration_slider = gr.Slider(minimum=5.0, maximum=20.0, step=5.0, value=20.0, label="Duration (Seconds)")
+                steps_slider = gr.Slider(minimum=2, maximum=50, step=1, value=4, label="Inference Steps (per chunk)")
+                cfg_slider = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.0, label="Guidance Scale")
+                cfg_slider_2 = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.0, label="Guidance Scale 2")
+                seed_input = gr.Number(label="Seed", value=42, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+
+            generate_button = gr.Button("GENERATE LONG VIDEO (20s)", variant="primary", size="lg")
+
+        with gr.Column(scale=1):
+            video_output = gr.Video(label="Generated Video")
+
+    ui_inputs = [
+        image_input, prompt_input, steps_slider, negative_prompt_input, 
+        duration_slider, cfg_slider, cfg_slider_2, seed_input, randomize_seed
+    ]
+
+    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
+
+    # --- BOTTOM ADVERTISEMENT BANNER ---
+    gr.HTML("""
+    <div style="background: linear-gradient(90deg, #4f46e5, #9333ea); color: white; padding: 15px; border-radius: 10px; text-align: center; margin-top: 20px; box-shadow: 0 4px 15px rgba(0,0,0,0.1);">
+        <div style="display: flex; align-items: center; justify-content: center; gap: 20px; flex-wrap: wrap;">
+            <div style="text-align: left;">
+                <h3 style="margin: 0; font-weight: bold; font-size: 18px;">✨ New: Dream Hub Pro (All-in-One)</h3>
+                <p style="margin: 5px 0 0 0; opacity: 0.9; font-size: 14px;">Access all your pro tools (Wan2.1, Qwen, Audio, Video Enhance) in one place!</p>
+            </div>
+            <a href="https://huggingface.co/spaces/dream2589632147/Dream-Hub-Pro" target="_blank" style="text-decoration: none;">
+                <button style="background-color: white; color: #4f46e5; border: none; padding: 10px 25px; border-radius: 25px; font-weight: bold; cursor: pointer; transition: all 0.2s; font-size: 15px; box-shadow: 0 2px 5px rgba(0,0,0,0.2);">
+                    🚀 Open Hub Pro Now
+                </button>
+            </a>
+        </div>
+    </div>
+    """)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch()