Update app.py

app.py

@@ -6,296 +6,157 @@ import gradio as gr
 import tempfile
 import spaces
 from huggingface_hub import hf_hub_download
-import logging
 import numpy as np
 from PIL import Image
-import random # Added for random seed generation
+import random
 
-# --- Global Model Loading & LoRA Handling ---
 MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 LORA_REPO_ID = "Kijai/WanVideo_comfy"
 LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# --- Model Loading ---
-logger.info(f"Loading Image Encoder for {MODEL_ID}...")
-image_encoder = CLIPVisionModel.from_pretrained(
-    MODEL_ID,
-    subfolder="image_encoder",
-    torch_dtype=torch.float32 # Using float32 for image encoder as sometimes bfloat16/float16 can be problematic
-)
-
-logger.info(f"Loading VAE for {MODEL_ID}...")
-vae = AutoencoderKLWan.from_pretrained(
-    MODEL_ID,
-    subfolder="vae",
-    torch_dtype=torch.float32 # Using float32 for VAE for precision
-)
-logger.info(f"Loading Pipeline {MODEL_ID}...")
+image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
+vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
 pipe = WanImageToVideoPipeline.from_pretrained(
-    MODEL_ID,
-    vae=vae,
-    image_encoder=image_encoder,
-    torch_dtype=torch.bfloat16 # Main pipeline can use bfloat16 for speed/memory
-)
-flow_shift = 8.0
-pipe.scheduler = UniPCMultistepScheduler.from_config(
-    pipe.scheduler.config, flow_shift=flow_shift
+    MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
 )
-logger.info("Moving pipeline to CUDA...")
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
 pipe.to("cuda")
 
-# --- LoRA Loading ---
-logger.info(f"Downloading LoRA {LORA_FILENAME} from {LORA_REPO_ID}...")
 causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
-
-logger.info("Loading LoRA weights...")
 pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
-logger.info("Setting LoRA adapter...")
-pipe.set_adapters(["causvid_lora"], adapter_weights=[1.0])
+pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
 
-# --- Constants for Dimension Calculation ---
 MOD_VALUE = 32
-MOD_VALUE_H = MOD_VALUE_W = MOD_VALUE
-
 DEFAULT_H_SLIDER_VALUE = 512
 DEFAULT_W_SLIDER_VALUE = 896
+NEW_FORMULA_MAX_AREA = 480.0 * 832.0 
 
-# New fixed max_area for the calculation formula
-NEW_FORMULA_MAX_AREA = float(480 * 832)
-
-SLIDER_MIN_H = 128
-SLIDER_MAX_H = 896
-SLIDER_MIN_W = 128
-SLIDER_MAX_W = 896
-
-# --- Constant for Seed ---
+SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
+SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
 MAX_SEED = np.iinfo(np.int32).max
 
+FIXED_FPS = 24
+MIN_FRAMES_MODEL = 8
+MAX_FRAMES_MODEL = 81 
+
 def _calculate_new_dimensions_wan(pil_image: Image.Image, mod_val: int, calculation_max_area: float,
                                  min_slider_h: int, max_slider_h: int,
                                  min_slider_w: int, max_slider_w: int,
                                  default_h: int, default_w: int) -> tuple[int, int]:
     orig_w, orig_h = pil_image.size
-
-    if orig_w <= 0 or orig_h <= 0: # Changed to <= 0 for robustness
-        logger.warning(f"Uploaded image has non-positive width or height ({orig_w}x{orig_h}). Using default slider dimensions.")
+    if orig_w <= 0 or orig_h <= 0:
         return default_h, default_w
 
     aspect_ratio = orig_h / orig_w
     
-    sqrt_h_term = np.sqrt(calculation_max_area * aspect_ratio)
-    sqrt_w_term = np.sqrt(calculation_max_area / aspect_ratio)
-
-    calc_h = round(sqrt_h_term) // mod_val * mod_val
-    calc_w = round(sqrt_w_term) // mod_val * mod_val
-
-    calc_h = mod_val if calc_h < mod_val else calc_h
-    calc_w = mod_val if calc_w < mod_val else calc_w
-
-    effective_min_h = min_slider_h
-    effective_min_w = min_slider_w
+    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
+    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
 
-    effective_max_h_from_slider = (max_slider_h // mod_val) * mod_val
-    effective_max_w_from_slider = (max_slider_w // mod_val) * mod_val
+    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))
     
-    new_h = int(np.clip(calc_h, effective_min_h, effective_max_h_from_slider))
-    new_w = int(np.clip(calc_w, effective_min_w, effective_max_w_from_slider))
-
-    logger.info(f"Auto-dim: Original {orig_w}x{orig_h} (AR: {aspect_ratio:.2f}). Max Area for calc: {calculation_max_area}.")
-    logger.info(f"Auto-dim: Sqrt terms HxW: {sqrt_h_term:.0f}x{sqrt_w_term:.0f}. Calculated (round(sqrt_term)//{mod_val}*{mod_val}): {calc_h}x{calc_w}.")
-    logger.info(f"Auto-dim: Clamped HxW: {new_h}x{new_w} (Effective H_range:[{effective_min_h}-{effective_max_h_from_slider}], Effective W_range:[{effective_min_w}-{effective_max_w_from_slider}]).")
-
     return new_h, new_w
 
 def handle_image_upload_for_dims_wan(uploaded_pil_image: Image.Image | None, current_h_val: int, current_w_val: int):
     if uploaded_pil_image is None:
-        logger.info("Image cleared. Resetting dimensions to default slider values.")
         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, # Use the globally defined max_area for the new formula
-            SLIDER_MIN_H, SLIDER_MAX_H,
-            SLIDER_MIN_W, SLIDER_MAX_W,
+            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:
-        logger.error(f"Error auto-adjusting H/W from image: {e}", exc_info=True)
-        # Fallback to default slider values on error, as in the original code
+        gr.Warning("Error attempting to calculate new dimensions")
         return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
 
-
-# --- Gradio Interface Function ---
 @spaces.GPU
 def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
                    height: int, width: int, duration_seconds: float,
                    guidance_scale: float, steps: int,
                    seed: int, randomize_seed: bool, 
                    progress=gr.Progress(track_tqdm=True)):
+    
     if input_image is None:
         raise gr.Error("Please upload an input image.")
 
-    # Constants for frame calculation
-    FIXED_FPS = 24
-    MIN_FRAMES_MODEL = 8
-    MAX_FRAMES_MODEL = 81
-
-    logger.info("Starting video generation...")
-    logger.info(f"  Input Image: Uploaded (Original size: {input_image.size if input_image else 'N/A'})")
-    logger.info(f"  Prompt: {prompt}")
-    logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
-    logger.info(f"  Target Output Height: {height}, Target Output Width: {width}")
-
-    target_height = int(height)
-    target_width = int(width)
-    guidance_scale_val = float(guidance_scale)
-    steps_val = int(steps)
-
-    num_frames_for_pipeline = int(round(duration_seconds * FIXED_FPS))
-    num_frames_for_pipeline = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, num_frames_for_pipeline))
-    if num_frames_for_pipeline < MIN_FRAMES_MODEL:
-        num_frames_for_pipeline = MIN_FRAMES_MODEL
-
-    logger.info(f"  Duration: {duration_seconds:.1f}s, Fixed FPS (conditioning & export): {FIXED_FPS}")
-    logger.info(f"  Calculated Num Frames: {num_frames_for_pipeline} (clamped to [{MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL}])")
-    logger.info(f"  Guidance Scale: {guidance_scale_val}, Steps: {steps_val}")
+    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
+    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
     
-    # Seed logic
-    current_seed = int(seed)
-    if randomize_seed:
-        current_seed = random.randint(0, MAX_SEED)
-    logger.info(f"  Initial Seed: {seed}, Randomize: {randomize_seed}, Using Seed: {current_seed}")
-
-
-    if target_height % MOD_VALUE_H != 0:
-        logger.warning(f"Height {target_height} is not a multiple of {MOD_VALUE_H}. Adjusting...")
-        target_height = (target_height // MOD_VALUE_H) * MOD_VALUE_H
-    if target_width % MOD_VALUE_W != 0:
-        logger.warning(f"Width {target_width} is not a multiple of {MOD_VALUE_W}. Adjusting...")
-        target_width = (target_width // MOD_VALUE_W) * MOD_VALUE_W
-
-    target_height = max(MOD_VALUE_H, target_height if target_height > 0 else MOD_VALUE_H)
-    target_width = max(MOD_VALUE_W, target_width if target_width > 0 else MOD_VALUE_W)
-
+    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)
 
-    resized_image = input_image.resize((target_width, target_height))
-    logger.info(f"  Input image resized to: {resized_image.size} for pipeline input.")
+    resized_image = input_image.resize((target_w, target_h))
 
     with torch.inference_mode():
         output_frames_list = pipe(
-            image=resized_image,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=target_height,
-            width=target_width,
-            num_frames=num_frames_for_pipeline,
-            guidance_scale=guidance_scale_val,
-            num_inference_steps=steps_val,
-            generator=torch.Generator(device="cuda").manual_seed(current_seed) # Use current_seed
+            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]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
-
     export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
-    logger.info(f"Video successfully generated and saved to {video_path}")
     return video_path
 
-# --- Gradio UI Definition ---
 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"
-penguin_image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
 
 with gr.Blocks() as demo:
-    gr.Markdown(f"""
-    # Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA
-    """)
+    gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA")
     with gr.Row():
         with gr.Column():
-            input_image_component = gr.Image(type="pil", label="Input Image (will be resized to target H/W)")
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
-            duration_seconds_input = gr.Slider(minimum=0.4, maximum=3.3, step=0.1, value=1.7, label="Duration (seconds)", info="The CausVid LoRA was trained on 24fps, Wan has 81 maximum frames limit, limiting the maximum to 3.3s")
+            input_image_component = gr.Image(type="pil", label="Input Image (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 (Optional)",
-                    value=default_negative_prompt,
-                    lines=3
-                )
-                # --- Added Seed Controls ---
-                seed_input = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=42, # Default seed value
-                    interactive=True
-                )
-                randomize_seed_checkbox = gr.Checkbox(
-                    label="Randomize seed",
-                    value=True, # Default to randomize
-                    interactive=True
-                )
-                # --- End of Added Seed Controls ---
+                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)
                 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=30, step=1, value=4, label="Inference Steps") 
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False) 
+                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)
 
             generate_button = gr.Button("Generate Video", variant="primary")
-
         with gr.Column():
-            video_output = gr.Video(label="Generated Video", interactive=False)
+            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(
+    
+    input_image_component.clear( 
         fn=handle_image_upload_for_dims_wan,
         inputs=[input_image_component, height_input, width_input],
         outputs=[height_input, width_input]
     )
-
-    inputs_for_click_and_examples = [
-        input_image_component,
-        prompt_input,
-        negative_prompt_input,
-        height_input,
-        width_input,
-        duration_seconds_input,
-        guidance_scale_input,
-        steps_slider,
-        seed_input, # Added seed_input
-        randomize_seed_checkbox # Added randomize_seed_checkbox
+    
+    ui_inputs = [
+        input_image_component, prompt_input, negative_prompt_input,
+        height_input, width_input, duration_seconds_input,
+        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
     ]
-
-    generate_button.click(
-        fn=generate_video,
-        inputs=inputs_for_click_and_examples,
-        outputs=video_output
-    )
+    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=video_output)
 
     gr.Examples(
-        examples=[
-            # Added seed (e.g., 42) and randomize_seed (e.g., True) to examples
-            ["peng.png", "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, 896, 512, 2, 1.0, 4, 42, False],
-            ["forg.jpg", "the frog jumps around", default_negative_prompt, 448, 832, 2, 1.0, 4, 123, False],
+        examples=[ 
+            ["peng.png", "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, 896, 512, 2.0, 1.0, 4, 42, False],
+            ["forg.jpg", "the frog jumps around", default_negative_prompt, 448, 832, 2.0, 1.0, 4, 123, False],
         ],
-        inputs=inputs_for_click_and_examples, 
-        outputs=video_output,
-        fn=generate_video,
-        cache_examples="lazy"
+        inputs=ui_inputs, outputs=video_output, fn=generate_video, cache_examples="lazy"
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(share=True, debug=True)
+    demo.queue().launch()