added cuda as optional

inference.py

@@ -55,7 +55,9 @@ def load_vae(vae_dir):
     vae = CausalVideoAutoencoder.from_config(vae_config)
     vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
     vae.load_state_dict(vae_state_dict)
-    return vae.cuda().to(torch.bfloat16)
+    if torch.cuda.is_available():
+        vae = vae.cuda()
+    return vae.to(torch.bfloat16)
 
 
 def load_unet(unet_dir):
@@ -65,7 +67,9 @@ def load_unet(unet_dir):
     transformer = Transformer3DModel.from_config(transformer_config)
     unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
     transformer.load_state_dict(unet_state_dict, strict=True)
-    return transformer.cuda()
+    if torch.cuda.is_available():
+        transformer = transformer.cuda()
+    return transformer
 
 
 def load_scheduler(scheduler_dir):
@@ -254,7 +258,9 @@ def main():
     patchifier = SymmetricPatchifier(patch_size=1)
     text_encoder = T5EncoderModel.from_pretrained(
         "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder"
-    ).to("cuda")
+    )
+    if torch.cuda.is_available():
+        text_encoder = text_encoder.to("cuda")
     tokenizer = T5Tokenizer.from_pretrained(
         "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
     )
@@ -272,7 +278,9 @@ def main():
         "vae": vae,
     }
 
-    pipeline = XoraVideoPipeline(**submodel_dict).to("cuda")
+    pipeline = XoraVideoPipeline(**submodel_dict)
+    if torch.cuda.is_available():
+        pipeline = pipeline.to("cuda")
 
     # Prepare input for the pipeline
     sample = {
@@ -286,8 +294,10 @@ def main():
     random.seed(args.seed)
     np.random.seed(args.seed)
     torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(args.seed)
+        
+    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else 'cpu').manual_seed(args.seed)
 
     images = pipeline(
         num_inference_steps=args.num_inference_steps,

xora/pipelines/pipeline_xora_video.py

@@ -1010,7 +1010,7 @@ class XoraVideoPipeline(DiffusionPipeline):
                     current_timestep = current_timestep * (1 - conditioning_mask)
                 # Choose the appropriate context manager based on `mixed_precision`
                 if mixed_precision:
-                    context_manager = torch.autocast("cuda", dtype=torch.bfloat16)
+                    context_manager = torch.autocast("cuda" if torch.cuda.is_available() else 'cpu', dtype=torch.bfloat16)
                 else:
                     context_manager = nullcontext()  # Dummy context manager