ui changes and bugfixes

official-templates/better-ai-launcher/app/app.py

@@ -121,25 +121,22 @@ def index():
             'status': status,
             'installed': dirs_ok,
             'install_status': install_status,
-            'is_bcomfy': app_name == 'bcomfy'  # Add this line
+            'is_bcomfy': app_name == 'bcomfy'
         }
+
     filebrowser_status = get_filebrowser_status()
     return render_template('index.html', 
-                           apps=app_configs, 
-                           app_status=app_status, 
-                           pod_id=RUNPOD_POD_ID, 
-                           RUNPOD_PUBLIC_IP=os.environ.get('RUNPOD_PUBLIC_IP'),
-                           RUNPOD_TCP_PORT_22=os.environ.get('RUNPOD_TCP_PORT_22'),
-
- # lutzapps - CHANGE #2 - allow localhost Url for unsecure "http" and "ws" WebSockets protocol,
- # according to LOCAL_DEBUG ENV var (used 3x in "index.html" changes)
-                           enable_unsecure_localhost=os.environ.get('LOCAL_DEBUG'),
-
-                           settings=settings,
-                           current_auth_method=current_auth_method,
-                           ssh_password=ssh_password,
-                           ssh_password_status=ssh_password_status,
-                           filebrowser_status=filebrowser_status)
+                         apps=app_configs, 
+                         app_status=app_status, 
+                         pod_id=RUNPOD_POD_ID, 
+                         RUNPOD_PUBLIC_IP=os.environ.get('RUNPOD_PUBLIC_IP'),
+                         RUNPOD_TCP_PORT_22=os.environ.get('RUNPOD_TCP_PORT_22'),
+                         enable_unsecure_localhost=os.environ.get('LOCAL_DEBUG'),
+                         settings=settings,
+                         current_auth_method=current_auth_method,
+                         ssh_password=ssh_password,
+                         ssh_password_status=ssh_password_status,
+                         filebrowser_status=filebrowser_status)
 
 @app.route('/start/<app_name>')
 def start_app(app_name):
@@ -520,35 +517,52 @@ def get_model_types_route():
 
 @app.route('/download_model', methods=['POST'])
 def download_model_route():
-    url = request.json.get('url')
-    model_name = request.json.get('model_name')
-    model_type = request.json.get('model_type')
-    civitai_token = request.json.get('civitai_token') or load_civitai_token()
-    hf_token = request.json.get('hf_token') or load_huggingface_token() # lutzapps - added HF_TOKEN ENV var support
-    version_id = request.json.get('version_id')
-    file_index = request.json.get('file_index')
-
-    is_civitai, _, _, _ = check_civitai_url(url)
-    is_huggingface, _, _, _, _ = check_huggingface_url(url) # TODO: double call
-
-    if not (is_civitai or is_huggingface):
-        return jsonify({'status': 'error', 'message': 'Unsupported URL. Please use Civitai or Hugging Face URLs.'}), 400
-
-    if is_civitai and not civitai_token:
-        return jsonify({'status': 'error', 'message': 'Civitai token is required for downloading from Civitai.'}), 400
-
     try:
-        success, message = download_model(url, model_name, model_type, civitai_token, hf_token, version_id, file_index)
-        if success:
-            if isinstance(message, dict) and 'choice_required' in message:
-                return jsonify({'status': 'choice_required', 'data': message['choice_required']})
-            return jsonify({'status': 'success', 'message': message})
-        else:
-            return jsonify({'status': 'error', 'message': message}), 400
+        data = request.json
+        url = data.get('url')
+        model_name = data.get('model_name')
+        model_type = data.get('model_type')
+        civitai_token = data.get('civitai_token')
+        hf_token = data.get('hf_token')
+        version_id = data.get('version_id')
+        file_index = data.get('file_index')
+
+        # If no token provided in request, try to read from file
+        if not civitai_token:
+            try:
+                if os.path.exists('/workspace/.civitai_token'):
+                    with open('/workspace/.civitai_token', 'r') as f:
+                        token_data = json.load(f)
+                        civitai_token = token_data.get('token')
+            except Exception as e:
+                app.logger.error(f"Error reading token file: {str(e)}")
+
+        is_civitai, _, _, _ = check_civitai_url(url)
+        is_huggingface, _, _, _, _ = check_huggingface_url(url)
+
+        if not (is_civitai or is_huggingface):
+            return jsonify({'status': 'error', 'message': 'Unsupported URL. Please use Civitai or Hugging Face URLs.'}), 400
+
+        if is_civitai and not civitai_token:
+            return jsonify({'status': 'error', 'message': 'Civitai token is required for downloading from Civitai.'}), 400
+
+        try:
+            success, message = download_model(url, model_name, model_type, civitai_token, hf_token, version_id, file_index)
+            if success:
+                if isinstance(message, dict) and 'choice_required' in message:
+                    return jsonify({'status': 'choice_required', 'data': message['choice_required']})
+                return jsonify({'status': 'success', 'message': message})
+            else:
+                return jsonify({'status': 'error', 'message': message}), 400
+        except Exception as e:
+            error_message = f"Model download error: {str(e)}\n{traceback.format_exc()}"
+            app.logger.error(error_message)
+            return jsonify({'status': 'error', 'message': error_message}), 500
+
     except Exception as e:
-        error_message = f"Model download error: {str(e)}\n{traceback.format_exc()}"
+        error_message = f"Error processing request: {str(e)}\n{traceback.format_exc()}"
         app.logger.error(error_message)
-        return jsonify({'status': 'error', 'message': error_message}), 500
+        return jsonify({'status': 'error', 'message': error_message}), 400
 
 @app.route('/get_model_folders')
 def get_model_folders():

official-templates/better-ai-launcher/app/utils/app_utils.py

@@ -14,7 +14,7 @@ import time
 import datetime
 import shutil
 from utils.app_configs import (app_configs, DEBUG_SETTINGS, pretty_dict, init_app_configs, init_debug_settings, write_debug_setting, ensure_kohya_local_venv_is_symlinked)
-from utils.model_utils import (get_sha256_hash_from_file)
+from utils.model_utils import (get_sha256_hash_from_file, download_civitai_model)
 
 INSTALL_STATUS_FILE = '/tmp/install_status.json'
 
@@ -701,6 +701,11 @@ def download_and_unpack_venv_v2(app_name:str, app_configs:dict, send_websocket_m
         send_websocket_message('install_progress', {'app_name': app_name, 'percentage': 100, 'stage': 'Unpacking Complete'})
         send_websocket_message('install_log', {'app_name': app_name, 'log': 'Unpacking complete. Proceeding to clone repository...'})
 
+        # Get the app config from app_configs
+        app_config = app_configs.get(app_name)
+        if not app_config:
+            return False, f"App '{app_name}' not found in configurations."
+
         # Clone the repository
         success, message = clone_application(app_config, send_websocket_message)
         if not success:
@@ -1074,3 +1079,71 @@ def install_app(app_name:str, app_configs:dict, send_websocket_message) -> tuple
         return success, message
     else:
         return False, f"Unknown app: {app_name}"
+
+def download_civitai_model(url, model_name, model_type, civitai_token=None, version_id=None, file_index=None):
+    try:
+        model_id = extract_model_id(url)
+        if not model_id:
+            raise ValueError("Could not extract model ID from URL")
+
+        # Use provided token or try to read from file
+        if not civitai_token:
+            try:
+                if os.path.exists('/workspace/.civitai_token'):
+                    with open('/workspace/.civitai_token', 'r') as f:
+                        token_data = json.load(f)
+                        civitai_token = token_data.get('token')
+            except Exception as e:
+                print(f"Error reading token file: {str(e)}")
+
+        headers = {"Authorization": f"Bearer {civitai_token}"} if civitai_token else {}
+        
+        # Get model info
+        model_info_url = f"https://civitai.com/api/v1/models/{model_id}"
+        response = requests.get(model_info_url, headers=headers)
+        if response.status_code != 200:
+            raise Exception(f"Failed to get model info: {response.text}")
+        
+        model_info = response.json()
+        versions = model_info.get('modelVersions', [])
+        
+        if not versions:
+            raise Exception("No versions found for this model")
+
+        # Handle version selection - return early if selection needed
+        if version_id is None and len(versions) > 1:
+            return {
+                'type': 'version',
+                'versions': versions,
+                'civitai_token': civitai_token
+            }
+        
+        # Get the selected version
+        selected_version = versions[0] if version_id is None else next(
+            (v for v in versions if str(v['id']) == str(version_id)), None)
+        
+        if not selected_version:
+            raise Exception("Specified version not found")
+
+        # Handle file selection - return early if selection needed
+        files = selected_version.get('files', [])
+        if not files:
+            raise Exception("No files found for this version")
+
+        if file_index is None and len(files) > 1:
+            return {
+                'type': 'file',
+                'files': files,
+                'version_id': selected_version['id'],
+                'civitai_token': civitai_token
+            }
+
+        # Select the file
+        selected_file = files[file_index if file_index is not None else 0]
+        if not selected_file:
+            raise Exception("No file found")
+
+        return selected_file
+
+    except Exception as e:
+        raise Exception(f"Exception downloading from CivitAI: {str(e)}")