fix

app.py

@@ -1,4 +1,5 @@
 import os
+import spaces
 import uuid
 from pathlib import Path
 import argparse
@@ -13,11 +14,11 @@ from simple_inference import load_model, inference_on_image
 # 1. Load model
 # -----------------------
 args = argparse.Namespace()
-args.blur2vid_hf_repo_path = "tedlasai/learn2refocus"
+args.learn2refocus_hf_repo_path = "tedlasai/learn2refocus"
 args.pretrained_model_path = "stabilityai/stable-video-diffusion-img2vid"
 args.seed = 0
 
-pipe, model_config = load_model(args)
+pipe, device = load_model(args)
 
 OUTPUT_DIR = Path("/tmp/output_stacks")
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
@@ -39,13 +40,13 @@ def generate_vstack_from_image(image: Image.Image, input_focal_position: int, nu
     args.num_inference_steps = num_inference_steps
 
     video_id = uuid.uuid4().hex
-    output_path = OUTPUT_DIR / f"{video_id}.mp4"
 
     args.device = "cuda"
 
     pipe.to(args.device)
-    processed_image, video = inference_on_image(pipe, image, interval_key, model_config, args)
-    export_to_video(video, output_path, fps=20)
+    batch = convert_to_batch(args.image_path, input_focal_position=input_focal_position)
+    output_frames, focal_stack_num = inference_on_image(args, batch, pipeline, device)
+    write_output(OUTPUT_DIR, output_frames, focal_stack_num, batch['icc_profile'])
 
     if not os.path.exists(output_path):
         raise gr.Error("Video generation failed: output file not found.")

requirements.txt

@@ -19,4 +19,4 @@ moviepy>=1.0.3
 pillow==9.5.0
 denku==0.0.51
 controlnet-aux==0.0.9
-gradio>=4.44.0
+gradio>=4.44.0

simplified_inference.py → simple_inference.py

@@ -27,9 +27,9 @@ from tqdm.auto import tqdm
 from transformers import CLIPVisionModelWithProjection
 from diffusers import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
 from diffusers.utils import check_min_version
-from simplified_pipeline import StableVideoDiffusionPipeline
-import videoio
+from simple_pipeline import StableVideoDiffusionPipeline
 from PIL import Image
+from diffusers.utils import export_to_video
 
 
 import argparse
@@ -174,12 +174,11 @@ def inference_on_image(args, batch, pipeline, device):
         num_inference_steps=args.num_inference_steps,
     )
     video_frames = svd_output.frames[0]
-
-
     video_frames_normalized = video_frames*0.5 + 0.5
     video_frames_normalized = torch.clamp(video_frames_normalized,0,1)
     video_frames_normalized = video_frames_normalized.permute(1,0,2,3)
-    video_frames_normalized = torch.nn.functional.interpolate(video_frames_normalized, ((pixel_values.shape[2]//2)*2, (pixel_values.shape[3]//2)*2), mode='bilinear')
+    video_frames_normalized = torch.nn.functional.interpolate(video_frames_normalized, ((pixel_values.shape[3]//2)*2, (pixel_values.shape[4]//2)*2), mode='bilinear')
+
 
     return video_frames_normalized, focal_stack_num
     # run inference
@@ -189,10 +188,8 @@ def write_output(output_dir, frames, focal_stack_num, icc_profile):
     print("Validation images will be saved to ", output_dir)
     os.makedirs(output_dir, exist_ok=True)
 
-    videoio.videosave(os.path.join(
-        output_dir,
-        f"stack.mp4",
-    ), frames.permute(0,2,3,1).cpu().numpy(), fps=5)
+    print("Frames shape: ", frames.shape)
+    export_to_video(frames.permute(0,2,3,1).cpu().numpy(), os.path.join(output_dir, "stack.mp4"), fps=20)
 
     #save images
     for i in range(9):