sampler.py

import base64
import json
from handler import EndpointHandler
from PIL import Image
import io
import os
import subprocess

# Initialize the handler
# base_dir = os.path.dirname(os.path.abspath(__file__))
# video_root = os.path.join(base_dir, "dw_poses_videos")
# extract_pose_path = os.path.join(base_dir, 'extract_dwpose_from_vid.py')
# command = f'python3 {extract_pose_path} --video_root {video_root}'

# # Run the command with shell=True
# result = subprocess.run(command, shell=True, capture_output=True, text=True)
# if result.returncode != 0:
#     raise RuntimeError(f"Error running extract_dwpose_from_vid.py: {result.stderr}")

handler = EndpointHandler()

# Define sample inputs
inputs = {
    "inputs": {
        "ref_image_url": "https://media.discordapp.net/attachments/1237667074210267217/1247398713228263494/image.jpg?ex=665fe214&is=665e9094&hm=ba3863d8149e4e43c6ec271b09fcd4027b6a58cc65b7a658797ab9fbc74b4808&=&format=webp&width=896&height=1194",
        "video_url": "https://firebasestorage.googleapis.com/v0/b/quiz-app-edffe.appspot.com/o/tell_my_bf_shortened_pose.mp4?alt=media&token=d3b7cb47-6fbb-435c-94a7-2cbd62705efb",
        "length": 48,
        "num_inference_steps": 15,
        "cfg": 3.5,
        "seed": -1,
        "firebase_doc_id" : "0gkm91Ao0GsmQ2WdLP9c"
    }
}

# Simulate an inference call
output = handler(inputs)

# # Decode the base64 video output
video_base64 = output.get("video", "")
video_bytes = base64.b64decode(video_base64)

# Save the video to a file
with open("output_video.mp4", "wb") as video_file:
    video_file.write(video_bytes)

print("Inference completed. Output video saved as output_video.mp4")