load models only if cuda available

app.py

@@ -21,29 +21,31 @@ from huggingface_hub import hf_hub_download
 # Ensure 'checkpoint' directory exists
 os.makedirs("checkpoints", exist_ok=True)
 
-# Download LoRA weights
-hf_hub_download(
-    repo_id="wenqsun/DimensionX",
-    filename="orbit_left_lora_weights.safetensors",
-    local_dir="checkpoints"
-)
+if not is_shared_ui and is_gpu_associated:
 
-hf_hub_download(
-    repo_id="wenqsun/DimensionX",
-    filename="orbit_up_lora_weights.safetensors",
-    local_dir="checkpoints"
-)
+    # Download LoRA weights
+    hf_hub_download(
+        repo_id="wenqsun/DimensionX",
+        filename="orbit_left_lora_weights.safetensors",
+        local_dir="checkpoints"
+    )
+
+    hf_hub_download(
+        repo_id="wenqsun/DimensionX",
+        filename="orbit_up_lora_weights.safetensors",
+        local_dir="checkpoints"
+    )
 
-# Load models in the global scope
-model_id = "THUDM/CogVideoX-5b-I2V"
-transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16).to("cpu")
-text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16).to("cpu")
-vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16).to("cpu")
-tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
-pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
+    # Load models in the global scope
+    model_id = "THUDM/CogVideoX-5b-I2V"
+    transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16).to("cpu")
+    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16).to("cpu")
+    vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16).to("cpu")
+    tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
+    pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
 
-# Add this near the top after imports
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+    # Add this near the top after imports
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
 def calculate_resize_dimensions(width, height, max_width=1024):
     """Calculate new dimensions maintaining aspect ratio"""