Update app.py

app.py

@@ -67,13 +67,18 @@ HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Check if running on Zero GPU (compile incompatible with Zero GPU)
 IS_ZEROGPU = os.environ.get("SPACES_ZERO_GPU") is not None
 
+# Определяем доступное устройство
+def get_available_device():
+    """Определяет доступное устройство (GPU или CPU)"""
+    if torch.cuda.is_available() and not IS_ZEROGPU:
+        return 'cuda'
+    else:
+        return 'cpu'
 
 # print("FISH_AE_DTYPE:", FISH_AE_DTYPE)
 # print("IS_ZEROGPU:", IS_ZEROGPU)
-# if IS_ZEROGPU:
-#     print("Running on Zero GPU - model compilation disabled")
-# else:
-#     print("Not on Zero GPU - model compilation available")
+# device = get_available_device()
+# print(f"Using device: {device}")
 
 def _safe_members(tf, prefix):
     if not prefix.endswith('/'):
@@ -170,11 +175,14 @@ def load_models():
     """Lazy load models on first use (required for Zero GPU)."""
     global model, model_compiled, fish_ae, pca_state, silentcipher_model
     if model is None:
-        # print("Loading models from HuggingFace...")
-        model = load_model_from_hf(dtype=MODEL_DTYPE, compile=False, token=HF_TOKEN)
-        fish_ae = load_fish_ae_from_hf(compile=(COMPILE_FISH_IF_NOT_ON_ZERO_GPU and not IS_ZEROGPU), dtype=FISH_AE_DTYPE, token=HF_TOKEN)
+        # Определяем доступное устройство
+        device = get_available_device()
+        # print(f"Loading models to device: {device}")
+        
+        model = load_model_from_hf(dtype=MODEL_DTYPE, compile=False, token=HF_TOKEN, device=device)
+        fish_ae = load_fish_ae_from_hf(compile=(COMPILE_FISH_IF_NOT_ON_ZERO_GPU and not IS_ZEROGPU), dtype=FISH_AE_DTYPE, token=HF_TOKEN, device=device)
 
-        pca_state = load_pca_state_from_hf(token=HF_TOKEN)
+        pca_state = load_pca_state_from_hf(token=HF_TOKEN, device=device)
         
         # Load silentcipher model if enabled
         if USE_SILENTCIPHER:
@@ -188,10 +196,6 @@ def load_models():
                 print("Continuing without watermarking...")
         
         # print("Models loaded successfully!")
-        # if not IS_ZEROGPU:
-        #     print("Note: model_compiled will be created when you check 'Compile Model'")
-
-
 
 def compile_model(should_compile):
     """Compile the model for faster inference."""
@@ -226,7 +230,6 @@ def do_compile():
     
     try:
         # Load models first if not already loaded (needed for compilation)
-        # Since Zero GPU can't compile, we can safely load eagerly here
         load_models()
         
         # print("Compiling model... This will take 1-3 minutes on first run.")
@@ -434,7 +437,9 @@ def generate_audio(
     
     # Load speaker audio if provided
     if speaker_audio_path is not None:
-        speaker_audio = load_audio(speaker_audio_path).cuda()
+        speaker_audio = load_audio(speaker_audio_path)
+        # Move to the same device as the model
+        speaker_audio = speaker_audio.to(active_model.device)
     else:
         speaker_audio = None
     
@@ -633,7 +638,9 @@ def generate_audio_simple(
     
     # Load speaker audio if provided
     if speaker_audio_path is not None:
-        speaker_audio = load_audio(speaker_audio_path).cuda()
+        speaker_audio = load_audio(speaker_audio_path)
+        # Move to the same device as the model
+        speaker_audio = speaker_audio.to(active_model.device)
     else:
         speaker_audio = None
     
@@ -2859,5 +2866,4 @@ if __name__ == "__main__":
     
     # Enable queue for better handling of concurrent requests on HF Spaces
     demo.queue(max_size=20)
-    demo.launch(allowed_paths=allowed_paths)
-
+    demo.launch(allowed_paths=allowed_paths)