Update README with project details and remove training script

README.md

@@ -1,76 +1,11 @@
-<div align="center">
-<h1>VMem: Consistent Video Scene Generation with Surfel-Indexed View Memory</h1>
-
-<a href="https://v-mem.github.io/"><img src="https://img.shields.io/badge/%F0%9F%8F%A0%20Project%20Page-gray.svg"></a>
-<a href="http://arxiv.org/abs/2503.14489"><img src="https://img.shields.io/badge/%F0%9F%93%84%20arXiv-2503.14489-B31B1B.svg"></a>
-<a href="https://huggingface.co/liguang0115/vmem"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange"></a>
-<a href="https://huggingface.co/spaces/stabilityai/stable-virtual-camera"><img src="https://img.shields.io/badge/%F0%9F%9A%80%20Gradio%20Demo-Huggingface-orange"></a>
-
-[Runjia Li](https://runjiali-rl.github.io/), [Philip Torr](https://www.robots.ox.ac.uk/~phst/), [Andrea Vedaldi](https://www.robots.ox.ac.uk/~vedaldi/), [Tomas Jakab](https://www.robots.ox.ac.uk/~tomj/)
-<br>
-<br>
-[University of Oxford](https://www.robots.ox.ac.uk/~vgg/)
-</div>
-
-<p align="center">
-  <img src="assets/demo_teaser.gif" width="100%" alt="Teaser" style="border-radius:10px;"/>
-</p>
-
-<!-- <p align="center" border-radius="10px">
-  <img src="assets/benchmark.png" width="100%" alt="teaser_page1"/>
-</p> -->
-
-# Overview
-
-`VMem` is a plug-and-play memory mechanism of image-set models for consistent scene generation.
-Existing methods either rely on inpainting with explicit geometry estimation, which suffers from inaccuracies, or use limited context windows in video-based approaches, leading to poor long-term coherence. To overcome these issues, we introduce Surfel Memory of Views (VMem), which anchors past views to surface elements (surfels) they observed. This enables conditioning novel view generation on the most relevant past views rather than just the most recent ones, enhancing long-term scene consistency while reducing computational cost.
-
-
-# :wrench: Installation
-
-```bash
-conda create -n vmem python=3.10
-conda activate vmem
-pip install -r requirements.txt
-```
-
-
-# :rocket: Usage
-
-You need to properly authenticate with Hugging Face to download our model weights. Once set up, our code will handle it automatically at your first run. You can authenticate by running
-
-```bash
-# This will prompt you to enter your Hugging Face credentials.
-huggingface-cli login
-```
-
-Once authenticated, go to our model card [here](https://huggingface.co/stabilityai/stable-virtual-camera) and enter your information for access.
-
-We provide a demo for you to interact with `VMem`. Simply run
-
-```bash
-python app.py
-```
-
-
-## :heart: Acknowledgement
-This work is built on top of [CUT3R](https://github.com/CUT3R/CUT3R), [DUSt3R](https://github.com/naver/dust3r) and [Stable Virtual Camera](https://github.com/stability-ai/stable-virtual-camera). We thank them for their great works.
-
-
-
-
-
-# :books: Citing
-
-If you find this repository useful, please consider giving a star :star: and citation.
-
-```
-@article{zhou2025stable,
-    title={Stable Virtual Camera: Generative View Synthesis with Diffusion Models},
-    author={Jensen (Jinghao) Zhou and Hang Gao and Vikram Voleti and Aaryaman Vasishta and Chun-Han Yao and Mark Boss and
-    Philip Torr and Christian Rupprecht and Varun Jampani
-    },
-    journal={arXiv preprint arXiv:2503.14489},
-    year={2025}
-}
-```
+---
+title: Stable Virtual Camera
+emoji: ⚡
+colorFrom: yellow
+colorTo: yellow
+sdk: gradio
+sdk_version: 5.33.0
+app_file: app.py
+pinned: false
+---
+- **Project Page**: [https://v-mem.github.io/](https://v-mem.github.io/)

train.py

@@ -1,263 +0,0 @@
-import argparse
-from datetime import datetime
-import random
-import os
-import time
-import multiprocessing
-
-# Set multiprocessing start method to 'spawn' to avoid CUDA initialization issues in forked processes
-multiprocessing.set_start_method('spawn', force=True)
-
-
-from tqdm.auto import tqdm  # Progress bar
-import numpy as np
-from omegaconf import OmegaConf
-
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from torch.optim.lr_scheduler import SequentialLR, LambdaLR, CosineAnnealingLR, ExponentialLR # Importing CosineAnnealingLR scheduler
-import torch.nn.functional as F
-
-
-
-from accelerate import Accelerator, DistributedDataParallelKwargs
-from accelerate.utils import set_seed  # Removed get_scheduler import
-
-from peft import get_peft_model, LoraConfig
-
-from modeling import VMemModel
-from modeling.modules.autoencoder import AutoEncoder
-from modeling.sampling import DDPMDiscretization, DiscreteDenoiser, create_samplers
-from modeling.modules.conditioner import CLIPConditioner
-
-from utils.training_utils import  DiffusionTrainer, load_pretrained_model
-from data.dataset import RealEstatePoseImageSevaDataset
-
-
-
-
-# set random seed for reproducibility
-torch.manual_seed(42)
-random.seed(42)
-np.random.seed(42)
-
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description='Train a model')
-    parser.add_argument('--config', type=str, default="", required=True, help='Path to the config file')
-    args = parser.parse_args()
-    return args
-
-
-def generate_current_datetime():
-    return datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-
-def prepare_model(unet, config):
-    assert isinstance(unet, VMemModel), "unet should be an instance of VMemModel"
-    if config.training.lora_flag:
-        target_modules = []
-        for name, param in unet.named_parameters():
-            # # if ("temporal" in name or "transformer" in name) and "norm" not in name:
-            print(name)
-            if ("transformer" in name or "emb" in name or "layers" in name) \
-                and "norm" not in name and "in_layers.0" not in name and "out_layers.0" not in name:
-                # print(name)
-                name = name.replace(".weight", "")
-                name = name.replace(".bias", "")
-                if name not in target_modules:
-                    target_modules.append(str(name))
-        
-        lora_config = LoraConfig(   
-            r=config.training.lora_r,
-            lora_alpha=config.training.lora_alpha,
-            target_modules=target_modules,
-            lora_dropout=config.training.lora_dropout,
-            # bias="none",
-        )
-        lora_config.target_modules = target_modules
-
-        unet = get_peft_model(unet, lora_config)
-        # for name, param in unet.named_parameters():
-        #     if "camera" in name or "control" in name or "context" in name or "epipolar" in name or "appearance" in name:
-        #         print(name)
-        #         param.requires_grad = True
-   
-        unet.print_trainable_parameters()
-    else:
-        for name, param in unet.named_parameters():
-            param.requires_grad = True
-         
-        print("trainable parameters percentage: ", np.sum([p.numel() for p in unet.parameters() if p.requires_grad])/np.sum([p.numel() for p in unet.parameters()]))
-    return unet
-
-
-    
- 
-def main():
-    args = parse_args()
-    config_path = args.config
-    config = OmegaConf.load(config_path)
-
-    # Load the configuration
-    num_epochs = config.training.num_epochs
-    batch_size = config.training.batch_size
-    learning_rate = config.training.learning_rate
-    gradient_accumulation_steps = config.training.gradient_accumulation_steps
-    num_workers = config.training.num_workers
-    warmup_epochs = config.training.warmup_epochs
-    max_grad_norm = config.training.max_grad_norm
-    validation_interval = config.training.validation_interval
-    visualization_flag = config.training.visualization_flag
-    visualize_every = config.training.visualize_every
-    random_seed = config.training.random_seed
-    save_flag = config.training.save_flag
-    use_wandb = config.training.use_wandb
-    samples_dir = config.training.samples_dir
-
-
-    
-    weights_save_dir = config.training.weights_save_dir
-    
-
-    resume = config.training.resume
-
-
-
-    exp_id = generate_current_datetime()
-    if visualization_flag:
-        run_visualization_dir = f"{samples_dir}/{exp_id}"
-        os.makedirs(run_visualization_dir, exist_ok=True)
-    else:
-        run_visualization_dir = None
-    if save_flag:
-        run_weights_save_dir = f"{weights_save_dir}/{exp_id}"
-        os.makedirs(run_weights_save_dir, exist_ok=True)
-    else:
-        run_weights_save_dir = None
-
-
-    accelerator = Accelerator(
-        mixed_precision="fp16",
-        gradient_accumulation_steps=gradient_accumulation_steps,
-        kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=False)],
-    )
-    num_gpus = accelerator.num_processes 
-
-    if random_seed is not None:
-        set_seed(random_seed, device_specific=True)
-    device = accelerator.device
-
-    
-
-    model = load_pretrained_model(cache_dir=config.model.cache_dir, device=device)
-
-
-    model = prepare_model(model, config)
-    if resume:
-        model.load_state_dict(torch.load(resume, map_location='cpu'), strict=False)
-        torch.cuda.empty_cache()
-    
-    # model = model.to(device)
-
-
-    # time.sleep(100*3600)
-
-      
-
-    train_dataset = RealEstatePoseImageSevaDataset(rgb_data_dir=config.dataset.realestate10k.rgb_data_dir, 
-                                                    meta_info_dir=config.dataset.realestate10k.meta_info_dir,
-                                                    num_sample_per_episode=config.dataset.realestate10k.num_sample_per_episode,
-                                                    mode='train')
-    val_dataset = RealEstatePoseImageSevaDataset(rgb_data_dir=config.dataset.realestate10k.rgb_data_dir, 
-                                                    meta_info_dir=config.dataset.realestate10k.meta_info_dir, 
-                                                    num_sample_per_episode=config.dataset.realestate10k.val_num_sample_per_episode,
-                                                    mode='test')
-
-        
-    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, multiprocessing_context='spawn')
-    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, multiprocessing_context='spawn')
-    
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=config.training.weight_decay)
-    train_steps_per_epoch = len(train_dataloader)
-    total_train_steps = num_epochs * train_steps_per_epoch 
-    warmup_steps = warmup_epochs * train_steps_per_epoch
-    
-    lr_scheduler = CosineAnnealingLR(
-        optimizer, T_max=total_train_steps - warmup_steps, eta_min=0
-    )
-    
-    # lr_scheduler = ExponentialLR(optimizer, gamma=gamma)
-    if warmup_epochs > 0:
-        def warmup_lambda(current_step):
-            return float(current_step) / float(max(1, warmup_steps))
-        warmup_scheduler = LambdaLR(optimizer, lr_lambda=warmup_lambda)
-
-
-        # Combine the schedulers using SequentialLR
-        lr_scheduler = SequentialLR(
-            optimizer, schedulers=[warmup_scheduler, lr_scheduler], milestones=[warmup_steps]
-        )
-    vae = AutoEncoder(chunk_size=1).to(device)
-    vae.eval()
-    conditioner = CLIPConditioner().to(device)
-    discretization = DDPMDiscretization()
-    denoiser = DiscreteDenoiser(discretization=discretization, num_idx=1000, device=device)
-    sampler = create_samplers(guider_types=config.training.guider_types,
-                              discretization=discretization,
-                              num_frames=config.model.num_frames,
-                              num_steps=config.training.inference_num_steps,
-                              cfg_min=config.training.cfg_min,
-                              device=device)
-
-
-    (model,
-    vae,
-    train_dataloader,
-    val_dataloader,
-    optimizer,
-    lr_scheduler) = accelerator.prepare(
-        model,
-        vae,
-        train_dataloader,
-        val_dataloader,
-        optimizer,
-        lr_scheduler,
-    )
-
-    
-    trainer = DiffusionTrainer(network=model,
-                               ae=vae,
-                               conditioner=conditioner,
-                               denoiser=denoiser,
-                               sampler=sampler,
-                               discretization=discretization,
-                               cfg=config.training.cfg,
-                               optimizer=optimizer,
-                               lr_scheduler=lr_scheduler,
-                               ema_decay=config.training.ema_decay,
-                               device=device,
-                               accelerator=accelerator,
-                               max_grad_norm=max_grad_norm,
-                               save_flag=save_flag,
-                               visualize_flag=visualization_flag)
-
-
-
-    trainer.train(train_dataloader, 
-                  num_epochs,
-                  unconditional_prob=config.training.uncond_prob,
-                  log_every=10, 
-                  validation_dataloader=val_dataloader, 
-                  validation_interval=validation_interval, 
-                  save_dir=run_weights_save_dir, 
-                  save_interval=config.training.save_every, 
-                  visualize_every=visualize_every, 
-                  visualize_dir=run_visualization_dir,
-                  use_wandb=use_wandb)
-
-
-if __name__ == "__main__":
-    main()
-