sync

.claude/settings.local.json

@@ -0,0 +1,10 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(rg:*)",
+      "Bash(find:*)",
+      "Bash(python3:*)"
+    ],
+    "deny": []
+  }
+}

app.py

@@ -536,11 +536,12 @@ with gr.Blocks(title="Wan2.1 1.3B LoRA Self-Forcing streaming demo") as demo:
                     
                     gr.Markdown("### ⚙️ Settings")
                     with gr.Row():
-                        seed = gr.Number(
-                            label="Seed", 
-                            value=-1, 
-                            info="Use -1 for random seed",
-                            precision=0
+                        seed = gr.Slider(
+                            label="Generation Seed (-1 for random)",
+                            minimum=-1,
+                            maximum=2147483647,  # 2^31 - 1
+                            step=1,
+                            value=-1
                         )
                         fps = gr.Slider(
                             label="Playback FPS", 

example_of_using_lora.py

@@ -0,0 +1,571 @@
+import gradio as gr
+import tempfile
+import random
+import json
+import os
+import shutil
+import hashlib
+import uuid
+from pathlib import Path
+import time
+import logging
+import torch
+import numpy as np
+from typing import Dict, Any, List, Optional, Tuple, Union
+from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import export_to_video
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Constants
+STORAGE_PATH = Path(os.getenv('STORAGE_PATH', './data'))
+LORA_PATH = STORAGE_PATH / "loras"
+OUTPUT_PATH = STORAGE_PATH / "output"
+
+MODEL_VERSION = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DEFAULT_PROMPT_PREFIX = ""
+
+# Create necessary directories
+STORAGE_PATH.mkdir(parents=True, exist_ok=True)
+LORA_PATH.mkdir(parents=True, exist_ok=True)
+OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
+
+# Global variables to track model state
+pipe = None
+current_lora_id = None
+
+def format_time(seconds: float) -> str:
+    """Format time duration in seconds to human readable string"""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    
+    parts = []
+    if hours > 0:
+        parts.append(f"{hours}h")
+    if minutes > 0:
+        parts.append(f"{minutes}m")
+    if secs > 0 or not parts:
+        parts.append(f"{secs}s")
+        
+    return " ".join(parts)
+
+def upload_lora_file(file: tempfile._TemporaryFileWrapper) -> Tuple[str, str]:
+    """Upload a LoRA file and return a hash-based ID for future reference
+    
+    Args:
+        file: Uploaded file object from Gradio
+        
+    Returns:
+        Tuple[str, str]: Hash-based ID for the stored file (returned twice for both outputs)
+    """
+    if file is None:
+        return "", ""
+        
+    try:
+        # Calculate SHA256 hash of the file
+        sha256_hash = hashlib.sha256()
+        with open(file.name, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                sha256_hash.update(chunk)
+        file_hash = sha256_hash.hexdigest()
+        
+        # Create destination path using hash
+        dest_path = LORA_PATH / f"{file_hash}.safetensors"
+        
+        # Check if file already exists
+        if dest_path.exists():
+            logger.info("LoRA file already exists")
+            return file_hash, file_hash
+        
+        # Copy the file to the destination
+        shutil.copy(file.name, dest_path)
+        
+        logger.info(f"a new LoRA file has been uploaded")
+        return file_hash, file_hash
+    except Exception as e:
+        logger.error(f"Error uploading LoRA file: {e}")
+        raise gr.Error(f"Failed to upload LoRA file: {str(e)}")
+
+def get_lora_file_path(lora_id: Optional[str]) -> Optional[Path]:
+    """Get the path to a LoRA file from its hash-based ID
+    
+    Args:
+        lora_id: Hash-based ID of the stored LoRA file
+        
+    Returns:
+        Path: Path to the LoRA file if found, None otherwise
+    """
+    if not lora_id:
+        return None
+        
+    # Check if file exists
+    lora_path = LORA_PATH / f"{lora_id}.safetensors"
+    if lora_path.exists():
+        return lora_path
+    
+    return None
+
+def get_or_create_pipeline(
+    enable_cpu_offload: bool = True, 
+    flow_shift: float = 3.0
+) -> WanPipeline:
+    """Get existing pipeline or create a new one if necessary
+    
+    Args:
+        enable_cpu_offload: Whether to enable CPU offload
+        flow_shift: Flow shift parameter for scheduler
+        
+    Returns:
+        WanPipeline: The pipeline for generation
+    """
+    global pipe
+    
+    if pipe is None:
+        # Create a new pipeline
+        logger.info("Creating new pipeline")
+        
+        # Load VAE
+        vae = AutoencoderKLWan.from_pretrained(MODEL_VERSION, subfolder="vae", torch_dtype=torch.float32)
+        
+        # Load transformer
+        pipe = WanPipeline.from_pretrained(MODEL_VERSION, vae=vae, torch_dtype=torch.bfloat16)
+        
+        # Configure scheduler
+        pipe.scheduler = UniPCMultistepScheduler.from_config(
+            pipe.scheduler.config, 
+            flow_shift=flow_shift
+        )
+        
+        # Move to GPU
+        pipe.to("cuda")
+        
+        # Enable CPU offload if requested
+        if enable_cpu_offload:
+            logger.info("Enabling CPU offload")
+            pipe.enable_model_cpu_offload()
+    else:
+        # Update existing pipeline's scheduler if needed
+        if pipe.scheduler.config.flow_shift != flow_shift:
+            logger.info(f"Updating scheduler flow_shift from {pipe.scheduler.config.flow_shift} to {flow_shift}")
+            pipe.scheduler = UniPCMultistepScheduler.from_config(
+                pipe.scheduler.config, 
+                flow_shift=flow_shift
+            )
+    
+    return pipe
+
+def manage_lora_weights(pipe: WanPipeline, lora_id: Optional[str], lora_weight: float) -> Tuple[bool, Optional[Path]]:
+    """Manage LoRA weights, loading/unloading only when necessary
+    
+    Args:
+        pipe: The pipeline to manage LoRA weights for
+        lora_id: UUID of LoRA file to use
+        lora_weight: Weight of LoRA contribution
+        
+    Returns:
+        Tuple[bool, Optional[Path]]: (Is using LoRA, Path to LoRA file)
+    """
+    global current_lora_id
+    
+    # Determine if we should use LoRA
+    using_lora = lora_id is not None and lora_id.strip() != "" and lora_weight > 0
+    
+    # If not using LoRA but we have one loaded, unload it
+    if not using_lora and current_lora_id is not None:
+        logger.info(f"Unloading current LoRA with ID")
+        try:
+            # Unload current LoRA weights
+            pipe.unload_lora_weights()
+            current_lora_id = None
+        except Exception as e:
+            logger.error(f"Error unloading LoRA weights: {e}")
+        return False, None
+    
+    # If using LoRA, check if we need to change weights
+    if using_lora:
+        lora_path = get_lora_file_path(lora_id)
+        
+        if not lora_path:
+            # Log the event but continue with base model
+            logger.warning(f"LoRA file with ID {lora_id} not found. Using base model instead.")
+            
+            # If we had a LoRA loaded, unload it
+            if current_lora_id is not None:
+                logger.info(f"Unloading current LoRA")
+                try:
+                    pipe.unload_lora_weights()
+                except Exception as e:
+                    logger.error(f"Error unloading LoRA weights: {e}")
+                current_lora_id = None
+                
+            return False, None
+        
+        # If LoRA ID changed, update weights
+        if lora_id != current_lora_id:
+            # If we had a LoRA loaded, unload it first
+            if current_lora_id is not None:
+                logger.info(f"Unloading current LoRA")
+                try:
+                    pipe.unload_lora_weights()
+                except Exception as e:
+                    logger.error(f"Error unloading LoRA weights: {e}")
+            
+            # Load new LoRA weights
+            logger.info("Using a LoRA")
+            try:
+                pipe.load_lora_weights(lora_path, weight_name=str(lora_path), adapter_name="default")
+                current_lora_id = lora_id
+            except Exception as e:
+                logger.error(f"Error loading LoRA weights: {e}")
+                return False, None
+        else:
+            logger.info(f"Using currently loaded LoRA with ID")
+        
+        return True, lora_path
+    
+    return False, None
+    
+def generate_video(
+    prompt: str,
+    negative_prompt: str,
+    prompt_prefix: str,
+    width: int,
+    height: int,
+    num_frames: int,
+    guidance_scale: float,
+    flow_shift: float,
+    lora_id: Optional[str],
+    lora_weight: float,
+    inference_steps: int,
+    fps: int = 16,
+    seed: int = -1,
+    enable_cpu_offload: bool = True,
+    conditioning_image: Optional[str] = None,
+    progress=gr.Progress()
+) -> str:
+    """Generate a video using the Wan model with optional LoRA weights
+    
+    Args:
+        prompt: Text prompt for generation
+        negative_prompt: Negative text prompt
+        prompt_prefix: Prefix to add to all prompts
+        width: Output video width
+        height: Output video height
+        num_frames: Number of frames to generate
+        guidance_scale: Classifier-free guidance scale
+        flow_shift: Flow shift parameter for scheduler
+        lora_id: UUID of LoRA file to use
+        lora_weight: Weight of LoRA contribution
+        inference_steps: Number of inference steps
+        fps: Frames per second for output video
+        seed: Random seed (-1 for random)
+        enable_cpu_offload: Whether to enable CPU offload for VRAM optimization
+        conditioning_image: Path to conditioning image for image-to-video (not used in this app)
+        progress: Gradio progress callback
+        
+    Returns:
+        str: Video path
+    """
+    global pipe, current_lora_id  # Move the global declaration to the top of the function
+    
+    try:
+        # Progress 0-5%: Initialize and check inputs
+        progress(0.00, desc="Initializing generation")
+        
+        # Add prefix to prompt
+        progress(0.02, desc="Processing prompt")
+        if prompt_prefix and not prompt.startswith(prompt_prefix):
+            full_prompt = f"{prompt_prefix}{prompt}"
+        else:
+            full_prompt = prompt
+        
+        # Create correct num_frames (should be 8*k + 1)
+        adjusted_num_frames = ((num_frames - 1) // 8) * 8 + 1
+        if adjusted_num_frames != num_frames:
+            logger.info(f"Adjusted number of frames from {num_frames} to {adjusted_num_frames} to match model requirements")
+            num_frames = adjusted_num_frames
+        
+        # Set up random seed
+        progress(0.03, desc="Setting up random seed")
+        if seed == -1:
+            seed = random.randint(0, 2**32 - 1)
+            logger.info(f"Using randomly generated seed: {seed}")
+        
+        # Set random seeds for reproducibility
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        generator = torch.Generator(device="cuda")
+        generator = generator.manual_seed(seed)
+        
+        # Progress 5-25%: Get or create pipeline
+        progress(0.05, desc="Preparing model")
+        pipe = get_or_create_pipeline(enable_cpu_offload, flow_shift)
+        
+        # Progress 25-40%: Manage LoRA weights
+        progress(0.25, desc="Managing LoRA weights")
+        using_lora, lora_path = manage_lora_weights(pipe, lora_id, lora_weight)
+        
+        # Create temporary file for the output
+        with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
+            output_path = temp_file.name
+        
+        # Progress 40-90%: Generate the video
+        progress(0.40, desc="Starting video generation")
+        
+        # Set up timing for generation
+        start_time = torch.cuda.Event(enable_timing=True)
+        end_time = torch.cuda.Event(enable_timing=True)
+        
+        start_time.record()
+        # Update progress once before generation starts
+        progress(0.45, desc="Running diffusion process")
+        
+        # Generate the video without callback
+        output = pipe(
+            prompt=full_prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            guidance_scale=guidance_scale,
+            num_inference_steps=inference_steps,
+            generator=generator,
+            # noo! don't do this!
+            # we will implement the lora weight / scale later
+            #cross_attention_kwargs={"scale": lora_weight} if using_lora else None
+        ).frames[0]
+        
+        # Update progress after generation completes
+        progress(0.90, desc="Generation complete")
+        
+        end_time.record()
+        torch.cuda.synchronize()
+        generation_time = start_time.elapsed_time(end_time) / 1000  # Convert to seconds
+        
+        logger.info(f"Video generation completed in {format_time(generation_time)}")
+        
+        # Progress 90-95%: Export video
+        progress(0.90, desc="Exporting video")
+        export_to_video(output, output_path, fps=fps)
+        
+        # Progress 95-100%: Save output and clean up
+        progress(0.95, desc="Saving video")
+        
+        # Save a copy to our output directory with UUID for potential future reference
+        output_id = str(uuid.uuid4())
+        saved_output_path = OUTPUT_PATH / f"{output_id}.mp4"
+        shutil.copy(output_path, saved_output_path)
+        logger.info(f"Saved video with ID: {output_id}")
+        
+        # No longer clear the pipeline since we're reusing it
+        # Just clean up local variables
+        progress(0.98, desc="Cleaning up resources")
+        
+        progress(1.0, desc="Generation complete")
+        
+        return output_path
+    
+    except Exception as e:
+        import traceback
+        error_msg = f"Error generating video: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        
+        # Clean up CUDA memory on error
+        if pipe is not None:
+            # Try to unload any LoRA weights on error
+            if current_lora_id is not None:
+                try:
+                    pipe.unload_lora_weights()
+                    current_lora_id = None
+                except:
+                    pass
+            
+            # Release the pipeline on critical errors
+            try:
+                pipe = None
+                torch.cuda.empty_cache()
+            except:
+                pass
+        
+        # Re-raise as Gradio error for UI display
+        raise gr.Error(f"Error generating video: {str(e)}")
+
+# Create the Gradio app
+with gr.Blocks(title="Video Generation API") as app:
+    
+    with gr.Tabs():
+        # LoRA Upload Tab
+        with gr.TabItem("1️⃣ Upload LoRA"):
+            gr.Markdown("## Upload LoRA Weights")
+            gr.Markdown("Upload your custom LoRA weights file to use for generation. The file will be automatically stored and you'll receive a unique hash-based ID.")
+            
+            with gr.Row():
+                lora_file = gr.File(label="LoRA File (safetensors format)")
+                
+            with gr.Row():
+                lora_id_output = gr.Textbox(label="LoRA Hash ID (use this in the generation tab)", interactive=False)
+            
+            # This will be connected after all components are defined
+        
+        # Video Generation Tab
+        with gr.TabItem("2️⃣ Generate Video"):
+            
+            with gr.Row():
+                with gr.Column(scale=1):
+                    # Input parameters
+                    prompt = gr.Textbox(
+                        label="Prompt",
+                        placeholder="Enter your prompt here...",
+                        lines=3
+                    )
+                    
+                    negative_prompt = gr.Textbox(
+                        label="Negative Prompt",
+                        placeholder="Enter negative prompt here...",
+                        lines=3,
+                        value="worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background"
+                    )
+                    
+                    prompt_prefix = gr.Textbox(
+                        label="Prompt Prefix",
+                        placeholder="Prefix to add to all prompts",
+                        value=DEFAULT_PROMPT_PREFIX
+                    )
+                    
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=1280,
+                            step=8,
+                            value=1280
+                        )
+                        
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=720,
+                            step=8,
+                            value=720
+                        )
+                    
+                    with gr.Row():
+                        num_frames = gr.Slider(
+                            label="Number of Frames",
+                            minimum=9,
+                            maximum=257,
+                            step=8,
+                            value=49
+                        )
+                        
+                        fps = gr.Slider(
+                            label="FPS",
+                            minimum=1,
+                            maximum=60,
+                            step=1,
+                            value=16
+                        )
+                    
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance Scale",
+                            minimum=1.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=5.0
+                        )
+                        
+                        flow_shift = gr.Slider(
+                            label="Flow Shift",
+                            minimum=0.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=3.0
+                        )
+                    
+                    lora_id = gr.Textbox(
+                        label="LoRA ID (from upload tab)",
+                        placeholder="Enter your LoRA ID here...",
+                    )
+                    
+                    with gr.Row():
+                        lora_weight = gr.Slider(
+                            label="LoRA Weight",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.01,
+                            value=0.7
+                        )
+                        
+                        inference_steps = gr.Slider(
+                            label="Inference Steps",
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=30
+                        )
+                    
+                    seed = gr.Slider(
+                        label="Generation Seed (-1 for random)",
+                        minimum=-1,
+                        maximum=2147483647,  # 2^31 - 1
+                        step=1,
+                        value=-1
+                    )
+                    
+                    enable_cpu_offload = gr.Checkbox(
+                        label="Enable Model CPU Offload (for low-VRAM GPUs)",
+                        value=False
+                    )
+                    
+                    generate_btn = gr.Button(
+                        "Generate Video",
+                        variant="primary"
+                    )
+                
+                with gr.Column(scale=1):
+                    # Output component - just the video preview
+                    preview_video = gr.Video(
+                        label="Generated Video",
+                        interactive=False
+                    )
+            
+            # Connect the generate button
+            generate_btn.click(
+                fn=generate_video,
+                inputs=[
+                    prompt,
+                    negative_prompt,
+                    prompt_prefix,
+                    width,
+                    height,
+                    num_frames,
+                    guidance_scale,
+                    flow_shift,
+                    lora_id,
+                    lora_weight,
+                    inference_steps,
+                    fps,
+                    seed,
+                    enable_cpu_offload
+                ],
+                outputs=[
+                    preview_video
+                ]
+            )
+    
+    # Connect LoRA upload to both display fields
+    lora_file.change(
+        fn=upload_lora_file,
+        inputs=[lora_file],
+        outputs=[lora_id_output, lora_id]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    app.launch()