fix

app.py

@@ -61,9 +61,9 @@ def handle_image_selection(image_name, state):
     print(f"User selected image: {image_name}")
     return initialize_simulator(image_name, state)
 
-if __name__ == '__main__':
-    with gr.Blocks() as demo:
-        genie_instance = gr.State({
+@spaces.GPU
+def init_model():
+    return gr.State({
             'genie': GenieSimulator(
                 image_encoder_type='temporalvae',
                 image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
@@ -76,6 +76,10 @@ if __name__ == '__main__':
             )
         })
 
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        genie_instance = init_model()
+
         with gr.Row():
             image_selector = gr.Dropdown(
                 choices=available_images, value=available_images[0], label="Select an Image"

sim/simulator.py

@@ -137,6 +137,7 @@ class GenieSimulator(LearnedSimulator):
         allow_external_prompt: bool = False
     ):
         super().__init__()
+        device = "cuda"
 
         assert quantize == (image_encoder_type == "magvit"), \
             "Currently quantization if and only if magvit is the image encoder."
@@ -378,10 +379,11 @@ class GenieSimulator(LearnedSimulator):
             W //= self.quant_slice_size
             _, _, indices, _ = self.image_encoder.encode(image, flip=True)
             indices = einops.rearrange(indices, "(h w) -> h w", h=H, w=W)
-            indices = indices.to(torch.int32)
+            indices = indices.to(torch.int32).to(self.device)
             return indices
 
         else:
+            self.image_encoder = self.image_encoder.to(device=self.device)
             if self.image_encoder_type == "magvit":
                 latent = self.image_encoder.encode_without_quantize(image)
             elif self.image_encoder_type == "temporalvae":
@@ -391,7 +393,7 @@ class GenieSimulator(LearnedSimulator):
                 latent = einops.rearrange(latent, "b c h w -> b h w c")
             else:
                 pass
-            latent = latent.squeeze(0).to(torch.float32)
+            latent = latent.squeeze(0).to(torch.float32).to(self.device)
             return latent