Create upload_ui.py

upload_ui.py

@@ -0,0 +1,62 @@
+"""Uploads for hosted sessions; no access to server filesystem paths."""
+from pathlib import Path
+
+import streamlit as st
+
+from app_environment import UPLOAD_ROOT, new_session_directory, remove_session_directory
+
+
+MAX_FILE_BYTES = 1024 * 1024 * 1024
+MAX_TOTAL_BYTES = 2 * MAX_FILE_BYTES
+
+
+def store_uploads(files, tokenizer=None, root=UPLOAD_ROOT):
+    if not files:
+        raise ValueError('Choose a GGUF or safetensors file.')
+    names = [item.name for item in files]
+    if any(not name or '/' in name or '\\' in name or ':' in name or '\x00' in name or Path(name).name != name for name in names):
+        raise ValueError('Use plain file names for uploaded model files.')
+    if len(set(name.casefold() for name in names)) != len(names):
+        raise ValueError('Uploaded model files must have different names.')
+    if any(Path(name).suffix.lower() not in ('.gguf', '.safetensors') for name in names):
+        raise ValueError('Upload GGUF or safetensors weights only.')
+    if any(Path(name).suffix.lower() == '.gguf' for name in names) and len(files) != 1:
+        raise ValueError('Upload one GGUF file, or one or more safetensors shards.')
+    all_files = list(files) + ([tokenizer] if tokenizer is not None else [])
+    sizes = [len(item.getbuffer()) for item in all_files]
+    if any(size > MAX_FILE_BYTES for size in sizes) or sum(sizes) > MAX_TOTAL_BYTES:
+        raise ValueError('Uploads allow up to 1 GB per file and 2 GB in total. Use Hugging Face for larger models.')
+    directory = new_session_directory(root)
+    try:
+        for item in files:
+            with (directory / item.name).open('wb') as target:
+                target.write(item.getbuffer())
+        tokenizer_path = None
+        if tokenizer is not None:
+            tokenizer_path = directory / 'tokenizer.json'
+            with tokenizer_path.open('wb') as target:
+                target.write(tokenizer.getbuffer())
+        model_path = directory / names[0] if len(names) == 1 else directory
+        return model_path, tokenizer_path, directory
+    except Exception:
+        remove_session_directory(directory, root)
+        raise
+
+
+def render_upload_loader(activate_model):
+    st.caption('Upload a single GGUF file or safetensors weights. Up to 1 GB per file, 2 GB total.')
+    files = st.file_uploader('Model files', type=['gguf', 'safetensors'], accept_multiple_files=True, key='uploaded_models')
+    tokenizer = st.file_uploader('Matching tokenizer.json (optional)', type=['json'], key='uploaded_tokenizer')
+    st.caption('Uploads are processed by this Space in your session and removed when you clear the model. Restarting the Space also clears uploads.')
+    if st.button('Load uploaded model', key='load_upload', type='primary', disabled=not files, use_container_width=True):
+        directory = None
+        try:
+            with st.spinner('Reading uploaded model…'):
+                path, tokenizer_path, directory = store_uploads(files, tokenizer)
+                activate_model(path, tokenizer_path, 'Uploaded: ' + ', '.join(item.name for item in files))
+            st.session_state.upload_directory = directory
+            st.rerun()
+        except Exception as error:
+            if directory is not None:
+                remove_session_directory(directory, UPLOAD_ROOT)
+            st.error(f'Could not load upload: {error}')