()
| 172 | raise ValueError(f"Unsupported input: {path}") |
| 173 | |
| 174 | def init_pipeline(): |
| 175 | print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device())) |
| 176 | mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu") |
| 177 | mm.load_models([ |
| 178 | "./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors", |
| 179 | ]) |
| 180 | pipe = FlashVSRTinyPipeline.from_model_manager(mm, device="cuda") |
| 181 | pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to("cuda", dtype=torch.bfloat16) |
| 182 | LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt" |
| 183 | if os.path.exists(LQ_proj_in_path): |
| 184 | pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu"), strict=True) |
| 185 | pipe.denoising_model().LQ_proj_in.to('cuda') |
| 186 | |
| 187 | multi_scale_channels = [512, 256, 128, 128] |
| 188 | pipe.TCDecoder = build_tcdecoder(new_channels=multi_scale_channels, new_latent_channels=16+768) |
| 189 | mis = pipe.TCDecoder.load_state_dict(torch.load("./FlashVSR/TCDecoder.ckpt"), strict=False) |
| 190 | print(mis) |
| 191 | |
| 192 | pipe.to('cuda'); pipe.enable_vram_management(num_persistent_param_in_dit=None) |
| 193 | pipe.init_cross_kv(); pipe.load_models_to_device(["dit","vae"]) |
| 194 | return pipe |
| 195 | |
| 196 | def main(): |
| 197 | RESULT_ROOT = "./results" |
no test coverage detected