Spaces:
Running
Running
| import os, json, re | |
| from googleapiclient.discovery import build | |
| from googleapiclient.http import MediaFileUpload | |
| from google.oauth2 import service_account | |
| import base64 | |
| from PIL import Image | |
| from PIL import Image | |
| from io import BytesIO | |
| from vouchervision.general_utils import get_cfg_from_full_path | |
| def setup_streamlit_config(dir_home): | |
| # Define the directory path and filename | |
| dir_path = os.path.join(dir_home, ".streamlit") | |
| file_path = os.path.join(dir_path, "config.toml") | |
| # Check if directory exists, if not create it | |
| if not os.path.exists(dir_path): | |
| os.makedirs(dir_path) | |
| # Create or modify the file with the provided content | |
| config_content = f""" | |
| [theme] | |
| base = "dark" | |
| primaryColor = "#00ff00" | |
| [server] | |
| enableStaticServing = false | |
| runOnSave = true | |
| port = 8524 | |
| maxUploadSize = 5000 | |
| """ | |
| with open(file_path, "w") as f: | |
| f.write(config_content.strip()) | |
| def save_uploaded_file(directory, img_file, image=None): | |
| if not os.path.exists(directory): | |
| os.makedirs(directory) | |
| # Assuming the uploaded file is an image | |
| if image is None: | |
| with Image.open(img_file) as image: | |
| full_path = os.path.join(directory, img_file.name) | |
| image.save(full_path, "JPEG") | |
| # Return the full path of the saved image | |
| return full_path | |
| else: | |
| full_path = os.path.join(directory, img_file.name) | |
| image.save(full_path, "JPEG") | |
| return full_path | |
| def image_to_base64(img): | |
| buffered = BytesIO() | |
| img.save(buffered, format="JPEG") | |
| return base64.b64encode(buffered.getvalue()).decode() | |
| def check_prompt_yaml_filename(fname): | |
| # Check if the filename only contains letters, numbers, underscores, and dashes | |
| pattern = r'^[\w-]+$' | |
| # The \w matches any alphanumeric character and is equivalent to the character class [a-zA-Z0-9_]. | |
| # The hyphen - is literally matched. | |
| if re.match(pattern, fname): | |
| return True | |
| else: | |
| return False | |
| # Function to upload files to Google Drive | |
| def upload_to_drive(filepath, filename): | |
| # Parse the service account info from the environment variable | |
| creds_info = json.loads(os.environ.get('GDRIVE_API')) | |
| if creds_info: | |
| creds = service_account.Credentials.from_service_account_info( | |
| creds_info, scopes=["https://www.googleapis.com/auth/drive"] | |
| ) | |
| service = build('drive', 'v3', credentials=creds) | |
| # Get the folder ID from the environment variable | |
| folder_id = os.environ.get('GDRIVE') | |
| # st.info(f"{folder_id}") | |
| if folder_id: | |
| file_metadata = { | |
| 'name': filename, | |
| 'parents': [folder_id] | |
| } | |
| # st.info(f"{file_metadata}") | |
| media = MediaFileUpload(filepath, mimetype='application/x-yaml') | |
| service.files().create( | |
| body=file_metadata, | |
| media_body=media, | |
| fields='id' | |
| ).execute() |