Update app.py

app.py

@@ -53,20 +53,6 @@ import numpy as np
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# LoRA Storage Configuration
-STORAGE_PATH = Path(DATA_ROOT) / "storage"
-LORA_PATH = STORAGE_PATH / "loras"
-OUTPUT_PATH = STORAGE_PATH / "output"
-
-# 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 for LoRA management
-current_lora_id = None
-current_lora_path = None
-
 # --- Argument Parsing ---
 parser = argparse.ArgumentParser(description="Gradio Demo for Self-Forcing with Frame Streaming")
 parser.add_argument('--port', type=int, default=7860, help="Port to run the Gradio app on.")
@@ -129,89 +115,6 @@ if not APP_STATE["torch_compile_applied"] and ENABLE_TORCH_COMPILATION:
     APP_STATE["torch_compile_applied"] = True
     print("✅ torch.compile applied to transformer")
 
-def upload_lora_file(file: tempfile._TemporaryFileWrapper) -> Tuple[str, str]:
-    """Upload a LoRA file and return a hash-based ID for future reference"""
-    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():
-            print(f"LoRA file already exists!")
-            return file_hash, file_hash
-        
-        # Copy the file to the destination
-        shutil.copy(file.name, dest_path)
-        
-        print(f"LoRA file uploaded!")
-        return file_hash, file_hash
-    except Exception as e:
-        print(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"""
-    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 manage_lora_weights(lora_id: Optional[str], lora_weight: float) -> Tuple[bool, Optional[Path]]:
-    """Manage LoRA weights for the transformer model"""
-    global current_lora_id, current_lora_path
-    
-    # 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, clear it
-    if not using_lora and current_lora_id is not None:
-        print(f"Clearing current LoRA")
-        current_lora_id = None
-        current_lora_path = None
-        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:
-            print(f"A LoRA file with this ID was found. Using base model instead.")
-            
-            # If we had a LoRA loaded, clear it
-            if current_lora_id is not None:
-                print(f"Clearing current LoRA")
-                current_lora_id = None
-                current_lora_path = None
-                
-            return False, None
-        
-        # If LoRA ID changed, update
-        if lora_id != current_lora_id:
-            print(f"Loading LoRA..")
-            current_lora_id = lora_id
-            current_lora_path = lora_path
-        else:
-            print(f"Using a LoRA!")
-        
-        return True, lora_path
-    
-    return False, None
-
 def frames_to_ts_file(frames, filepath, fps = 15):
     """
     Convert frames directly to .ts file using PyAV.
@@ -327,7 +230,7 @@ pipeline = CausalInferencePipeline(
 pipeline.to(dtype=torch.float16).to(gpu)
 
 @torch.no_grad()
-def video_generation_handler_streaming(prompt, seed=42, fps=15, width=400, height=224, duration=5, lora_id=None, lora_weight=0.0):
+def video_generation_handler_streaming(prompt, seed=42, fps=15, width=400, height=224, duration=5):
     """
     Generator function that yields .ts video chunks using PyAV for streaming.
     """
@@ -336,13 +239,6 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15, width=400, heigh
     
     # print(f"🎬 Starting PyAV streaming: seed: {seed}, duration: {duration}s")
 
-    # Handle LoRA weights
-    using_lora, lora_path = manage_lora_weights(lora_id, lora_weight)
-    if using_lora:
-        print(f"🎨 Using LoRA with weight factor {lora_weight}")
-    else:
-        print("🎨 Using base model (no LoRA)")
-
     # Setup
     conditional_dict = text_encoder(text_prompts=[prompt])
     for key, value in conditional_dict.items():
@@ -504,133 +400,93 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15, width=400, heigh
     print(f"✅ PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
 # --- Gradio UI Layout ---
-with gr.Blocks(title="Wan2.1 1.3B LoRA Self-Forcing streaming demo") as demo:
-    gr.Markdown("# 🚀 Run Any LoRA in near real-time!")
-    gr.Markdown("Real-time video generation with distilled Wan2-1 1.3B and LoRA [[Model]](https://huggingface.co/gdhe17/Self-Forcing), [[Project page]](https://self-forcing.github.io), [[Paper]](https://huggingface.co/papers/2506.08009)")
+with gr.Blocks(title="Wan2.1 1.3B Self-Forcing streaming demo") as demo:
+    gr.Markdown("Real-time video generation with distilled Wan2-1 1.3B [[Model]](https://huggingface.co/gdhe17/Self-Forcing), [[Project page]](https://self-forcing.github.io), [[Paper]](https://huggingface.co/papers/2506.08009)")
     
-    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():
+       with gr.Column(scale=2):
+           with gr.Group():
+               prompt = gr.Textbox(
+                   label="Prompt", 
+                   placeholder="A stylish woman walks down a Tokyo street...", 
+                   lines=4,
+                   value=""
+               )
+           start_btn = gr.Button("🎬 Start Streaming", variant="primary", size="lg")
+           
+           gr.Markdown("### ⚙️ Settings")
+           with gr.Row():
+               seed = gr.Number(
+                   label="Seed", 
+                   value=-1, 
+                   info="Use -1 for random seed",
+                   precision=0
+               )
+               fps = gr.Slider(
+                   label="Playback FPS", 
+                   minimum=1, 
+                   maximum=30, 
+                   value=args.fps, 
+                   step=1,
+                   visible=False,
+                   info="Frames per second for playback"
+               )
+                   
             with gr.Row():
-                lora_file = gr.File(label="LoRA File (safetensors format)")
-                
+                duration = gr.Slider(
+                    label="Duration (seconds)",
+                    minimum=1,
+                    maximum=5,
+                    value=3,
+                    step=1,
+                    info="Video duration in seconds"
+                )
+            
             with gr.Row():
-                lora_id_output = gr.Textbox(label="LoRA Hash ID (use this in the generation tab)", interactive=False)
+                width = gr.Slider(
+                    label="Width", 
+                    minimum=224, 
+                    maximum=720, 
+                    value=400, 
+                    step=8,
+                    info="Video width in pixels (8px steps)"
+                )
+                height = gr.Slider(
+                    label="Height", 
+                    minimum=224, 
+                    maximum=720, 
+                    value=224, 
+                    step=8,
+                    info="Video height in pixels (8px steps)"
+                )
         
-        # Video Generation Tab
-        with gr.TabItem("2️⃣ Generate Video"):
-            with gr.Row():
-                with gr.Column(scale=2):
-                    with gr.Group():
-                        prompt = gr.Textbox(
-                            label="Prompt", 
-                            placeholder="A stylish woman walks down a Tokyo street...", 
-                            lines=4,
-                            value=""
-                        )
-
-                    start_btn = gr.Button("🎬 Start Streaming", variant="primary", size="lg")
-                    
-                    gr.Markdown("### ⚙️ Settings")
-                    with gr.Row():
-                        seed = gr.Number(
-                            label="Seed", 
-                            value=-1, 
-                            info="Use -1 for random seed",
-                            precision=0
-                        )
-                        fps = gr.Slider(
-                            label="Playback FPS", 
-                            minimum=1, 
-                            maximum=30, 
-                            value=args.fps, 
-                            step=1,
-                            visible=False,
-                            info="Frames per second for playback"
-                        )
-                    
-                    with gr.Row():
-                        duration = gr.Slider(
-                            label="Duration (seconds)",
-                            minimum=1,
-                            maximum=5,
-                            value=3,
-                            step=1,
-                            info="Video duration in seconds"
-                        )
-                    
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width", 
-                            minimum=224, 
-                            maximum=720, 
-                            value=400, 
-                            step=8,
-                            info="Video width in pixels (8px steps)"
-                        )
-                        height = gr.Slider(
-                            label="Height", 
-                            minimum=224, 
-                            maximum=720, 
-                            value=224, 
-                            step=8,
-                            info="Video height in pixels (8px steps)"
-                        )
-                    
-                    gr.Markdown("### 🎨 LoRA Settings")
-                    lora_id = gr.Textbox(
-                        label="LoRA ID (from upload tab)",
-                        placeholder="Enter your LoRA ID here...",
-                    )
-                    
-                    lora_weight = gr.Slider(
-                        label="LoRA Weight",
-                        minimum=0.0,
-                        maximum=1.0,
-                        step=0.01,
-                        value=1.0,
-                        info="Strength of LoRA influence"
-                    )
-                    
-                with gr.Column(scale=3):
-                    gr.Markdown("### 📺 Video Stream")
-
-                    streaming_video = gr.Video(
-                        label="Live Stream",
-                        streaming=True,
-                        loop=True,
-                        height=400,
-                        autoplay=True,
-                        show_label=False
-                    )
-                    
-                    status_display = gr.HTML(
-                        value=(
-                            "<div style='text-align: center; padding: 20px; color: #666; border: 1px dashed #ddd; border-radius: 8px;'>"
-                            "🎬 Ready to start streaming...<br>"
-                            "<small>Configure your prompt and click 'Start Streaming'</small>"
-                            "</div>"
-                        ),
-                        label="Generation Status"
-                    )
+        with gr.Column(scale=3):
+            gr.Markdown("### 📺 Video Stream")
+            streaming_video = gr.Video(
+                label="Live Stream",
+                streaming=True,
+                loop=True,
+                height=400,
+                autoplay=True,
+                show_label=False
+            )
+            
+            status_display = gr.HTML(
+                value=(
+                    "<div style='text-align: center; padding: 20px; color: #666; border: 1px dashed #ddd; border-radius: 8px;'>"
+                    "🎬 Ready to start streaming...<br>"
+                    "<small>Configure your prompt and click 'Start Streaming'</small>"
+                    "</div>"
+                ),
+                label="Generation Status"
+            )
 
     # Connect the generator to the streaming video
     start_btn.click(
         fn=video_generation_handler_streaming,
-        inputs=[prompt, seed, fps, width, height, duration, lora_id, lora_weight],
+        inputs=[prompt, seed, fps, width, height, duration],
         outputs=[streaming_video, status_display]
     )
-    
-    # Connect LoRA upload to both display fields
-    lora_file.change(
-        fn=upload_lora_file,
-        inputs=[lora_file],
-        outputs=[lora_id_output, lora_id]
-    )
-
 
 # --- Launch App ---
 if __name__ == "__main__":