adding directly to firebase

handler.py

@@ -31,6 +31,10 @@ from rembg import remove
 import onnxruntime as ort
 import shutil
 
+import firebase_admin
+from firebase_admin import credentials, storage, firestore
+import json
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 if device.type != 'cuda':
@@ -44,6 +48,19 @@ class EndpointHandler():
         if not os.path.exists(config_path):
             raise FileNotFoundError(f"The configuration file was not found at: {config_path}")
 
+        service_account_info = os.getenv("FIREBASE_ACCOUNT_INFO")
+
+        if not service_account_info:
+            raise ValueError("The FIREBASE_SERVICE_ACCOUNT environment variable is not set.")
+        service_account_info = service_account_info.replace('/\\n/g', '\n')
+
+        service_account_info_dict = json.loads(service_account_info)
+
+        cred = credentials.Certificate(service_account_info_dict)
+        firebase_admin.initialize_app(cred, {
+            'storageBucket': 'quiz-app-edffe.appspot.com'
+        })
+
         self.config = OmegaConf.load(config_path)
         self.weight_dtype = torch.float16
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -277,10 +294,10 @@ class EndpointHandler():
         num_inference_steps = inputs.get("num_inference_steps", 25)
         cfg = inputs.get("cfg", 3.5)
         seed = inputs.get("seed", 123)
+        firebase_doc_id = inputs.get("firebase_doc_id", "")
 
         base_dir = os.path.dirname(os.path.abspath(__file__))
 
-        # Create a unique temporary directory for this request
         with tempfile.TemporaryDirectory() as temp_dir:
             print(f"Temporary directory created at {temp_dir}")  # Debug statement
             video_root = os.path.join(temp_dir, "dw_poses_videos")
@@ -288,14 +305,10 @@ class EndpointHandler():
             downloaded_video_path = os.path.join(video_root, "downloaded_video.mp4")
             downloaded_image_path = os.path.join(video_root, "downloaded_image.jpg")
 
-            # Download the video from the URL
             self.download_file(video_url, downloaded_video_path)
-
-            # Download the reference image from the URL
             self.download_file(ref_image_url, downloaded_image_path)
             ref_image = Image.open(downloaded_image_path)
 
-            # Calculate new dimensions
             original_width, original_height = ref_image.size
             max_dimension = max(original_width, original_height)
             if max_dimension > 600:
@@ -306,41 +319,12 @@ class EndpointHandler():
                 width = original_width
                 height = original_height
 
-            # Remove the background from the reference image
             ref_image_no_bg = remove(ref_image)
             ref_image_no_bg_path = os.path.join(video_root, "ref_image_no_bg.png")
             ref_image_no_bg.save(ref_image_no_bg_path)
 
-            # pose_output_path = os.path.join(temp_dir, "pose_videos")
-
-            # print("we are number 1")
-            # # Run the extract_dwpose_from_vid.py script
-            # 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}")
-            # print("we are number 2")
-
-            # # Locate the extracted pose video
-            # save_dir = video_root + "_dwpose"
-            # print(f"Expected save directory: {save_dir}")  # Debug statement
-            # pose_video_path = os.path.join(save_dir, "downloaded_video.mp4")
-
-            # if not os.path.exists(pose_video_path):
-            #     print("Contents of the temporary directory:")
-            #     self.print_directory_contents(temp_dir)
-            #     raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}")
-
-            # Speed up the pose video by 4x
-            # sped_up_pose_video_path = os.path.join(temp_dir, "sped_up_pose_video.mp4")
-            # self.speed_up_video(pose_video_path, sped_up_pose_video_path, factor=4)
-
             torch.manual_seed(seed)
 
-            # temporary_pose_video_path = os.path.join(base_dir, 'dancing.mp4')
             pose_images = read_frames(downloaded_video_path)
             src_fps = get_fps(downloaded_video_path)
             
@@ -365,53 +349,36 @@ class EndpointHandler():
             animation_path = os.path.join(save_dir, "animation_output.mp4")
             save_videos_grid(video, animation_path, n_rows=1, fps=src_fps)
 
-            # Crop the face from the reference image and save it
             cropped_face_path = os.path.join(save_dir, "cropped_face.jpg")
             cropped_face = self._crop_face(ref_image_no_bg, save_path=cropped_face_path)
 
             torch.cuda.empty_cache()
 
-            # Perform face swapping
-            # self.print_directory_contents()
             swapped_face_video_path = os.path.join(save_dir, "swapped_face_output.mp4")
-            
-            # Subprocess call to facefusion for face swapping
-            # self.print_directory_contents()
             facefusion_script_path = os.path.join(base_dir, 'facefusion', 'core.py')
             swap_command = f'python3 {facefusion_script_path} --source {cropped_face_path} --target {animation_path} --output {swapped_face_video_path}'
             swap_result = subprocess.run(swap_command, shell=True, capture_output=True, text=True)
             if swap_result.returncode != 0:
                 raise RuntimeError(f"Error running face swap: {swap_result.stderr}")
 
-            # Slow down the produced video by 4x
-            # self.print_directory_contents(temp_dir)
-            # slowed_down_animation_path = os.path.join(save_dir, "slowed_down_animation_output.mp4")
-            # self.slow_down_video(animation_path, slowed_down_animation_path, factor=4)
-
-            # Clear CUDA cache before RIFE interpolation
             torch.cuda.empty_cache()
 
-
-            #remove background
-            # self.print_directory_contents()
-            # removed_background_output_path = os.path.join(save_dir, "removed_background_result.mp4")
-            # remove_background_script_path = os.path.join(base_dir, "rembg_video.py")
-            # remove_background_command = f'python3 {remove_background_script_path} {swapped_face_video_path} {removed_background_output_path}'
-            # print("Command is " + remove_background_command)
-            # remove_background_result = subprocess.run(remove_background_command, shell=True, capture_output=True, text=True)
-            # if remove_background_result.returncode != 0:
-            #     raise RuntimeError(f"Error running removing backgriund: {remove_background_result.stderr}")
-
-
-            # Perform RIFE interpolation
-            # self.print_directory_contents(temp_dir)
-            # rife_output_path = os.path.join(save_dir, "completed_result.mp4")
-            # self.run_rife_interpolation(swapped_face_video_path, rife_output_path, multi=2, scale=0.5)
-
-            # Encode the final video in base64
             with open(swapped_face_video_path, "rb") as video_file:
                 video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
 
-            torch.cuda.empty_cache()
+            # Upload video to Firebase Storage
+            bucket = storage.bucket()
+            blob = bucket.blob(f"videos/{firebase_doc_id}/swapped_face_output.mp4")
+            blob.upload_from_filename(swapped_face_video_path)
+            
+            # Make the file publicly accessible
+            blob.make_public()
+
+            video_url = blob.public_url
+
+            # Update Firestore document
+            db = firestore.client()
+            doc_ref = db.collection('danceResults').document(firebase_doc_id)
+            doc_ref.update({"videoURL": video_url})
 
             return {"video": video_base64}

requirements.txt

@@ -61,4 +61,7 @@ mediapipe==0.9.1.0
 filetype==1.2.0
 
 
-transparent-background
+transparent-background
+
+
+firebase-admin==6.0.1

sampler.py

@@ -24,10 +24,11 @@ inputs = {
     "inputs": {
         "ref_image_url": "https://media.discordapp.net/attachments/1237667074210267217/1247311431028703242/image.png?ex=665f90cb&is=665e3f4b&hm=89767149275a6ef40e39d4a73d77019ab30ab665512398816bc614cfecdddb34&=&format=webp&quality=lossless&width=590&height=1194",
         "video_url": "https://firebasestorage.googleapis.com/v0/b/quiz-app-edffe.appspot.com/o/creeper_shortneed_pose.mp4?alt=media&token=f87f6e8c-6ab2-445c-8b55-205f08f1f5d5",
-        "length": 96,
+        "length": 6,
         "num_inference_steps": 25,
         "cfg": 3.5,
-        "seed": 123
+        "seed": 123,
+        "firebase_doc_id" : "0gkm91Ao0GsmQ2WdLP9c"
     }
 }