Refactor model initialization and GPU handling in generate function

- Simplify model initialization logic
- Ensure model is initialized with correct device
- Remove redundant GPU logging
- Improve device selection for model loading

webui.py

@@ -50,14 +50,15 @@ def generate(text,
     """Generate audio from text."""
     
     global MODEL
-    model = MODEL
     
-    if model is None:
-        raise RuntimeError("Model not initialized. Please ensure the model is loaded before generating audio.")
+    # Initialize model if not already done
+    if MODEL is None:
+        MODEL = initialize_model(device="cuda" if torch.cuda.is_available() else "cpu")
+    
+    model = MODEL
     
     # if gpu available, move model to gpu
     if torch.cuda.is_available():
-        print("Moving model to GPU")
         model = model.to("cuda")
     
     with torch.no_grad():
@@ -117,9 +118,12 @@ def build_ui(model_dir, device=0):
     
     global MODEL
     
-    # Initialize model
+    # Initialize model with proper device handling
+    device = "cuda" if torch.cuda.is_available() and device != "cpu" else "cpu"
     if MODEL is None:
         MODEL = initialize_model(model_dir, device=device)
+        if device == "cuda":
+            MODEL = MODEL.to(device)
 
     # Define callback function for voice cloning
     def voice_clone(text, prompt_text, prompt_wav_upload, prompt_wav_record):