updating app

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 from PIL import Image
 from diffusers.utils import export_to_video
 
-from inference import load_model, inference_on_image
+from simple_inference import load_model, inference_on_image
 
 # -----------------------
 # 1. Load model

simplified_inference.py

@@ -149,19 +149,12 @@ def convert_to_batch(image, input_focal_position, sample_frames=9):
         return {"pixel_values": pixels, "focal_stack_num": focal_stack_num, "original_pixel_values": original_pixels, 'icc_profile': icc_profile, "name": name}
 
 
-def inference_on_image(args, batch, unet, image_encoder, vae, global_step, weight_dtype, device):
+def inference_on_image(args, batch, pipeline, device):
+
 
-    pipeline = StableVideoDiffusionPipeline.from_pretrained(
-        args.pretrained_model_path,
-        unet=unet,
-        image_encoder=image_encoder,
-        vae=vae,
-        torch_dtype=weight_dtype,
-    )
 
     pipeline.set_progress_bar_config(disable=True)
     num_frames = 9 
-    unet.eval()
 
     pixel_values = batch["pixel_values"].to(device)
     focal_stack_num = batch["focal_stack_num"]
@@ -209,19 +202,10 @@ def write_output(output_dir, frames, focal_stack_num, icc_profile):
             img.info['icc_profile'] = icc_profile
         img.save(os.path.join(output_dir, f"frame_{i}.png"))
 
-
-def main():
-    args = parse_args()
-
-    if args.seed is not None:
-        set_seed(args.seed)
-
-    if args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
+def load_model(args):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # inference-only modules
+        # inference-only modules
     image_encoder = CLIPVisionModelWithProjection.from_pretrained(
         args.pretrained_model_path, subfolder="image_encoder"
     )
@@ -237,12 +221,35 @@ def main():
     unet = UNetSpatioTemporalConditionModel.from_pretrained(
         args.learn2refocus_hf_repo_path, subfolder="checkpoint-200000/unet"
     ).to(device)
+    unet.eval(); image_encoder.eval(); vae.eval()
+
+
+    pipeline = StableVideoDiffusionPipeline.from_pretrained(
+        args.pretrained_model_path,
+        unet=unet,
+        image_encoder=image_encoder,
+        vae=vae,
+        torch_dtype=weight_dtype,
+    )
+    return pipeline, device
+
+
+def main():
+    args = parse_args()
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+
+    pipeline, device = load_model(args)
 
     batch = convert_to_batch(args.image_path, input_focal_position=6)
 
-    unet.eval(); image_encoder.eval(); vae.eval()
     with torch.no_grad():
-        output_frames, focal_stack_num = inference_on_image(args, batch, unet, image_encoder, vae, 0, weight_dtype, device)
+        output_frames, focal_stack_num = inference_on_image(args, batch, pipeline, device)
         val_save_dir = os.path.join(args.output_dir, "validation_images", batch['name'])
         write_output(val_save_dir, output_frames, focal_stack_num, batch['icc_profile'])