-
Notifications
You must be signed in to change notification settings - Fork 12
/
share_model.py
96 lines (81 loc) · 3.75 KB
/
share_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# get save_dir_id, save_dir = save_dir_id
def get_save_dir_id (save_dir):
save_dir_id = None
try:
path_parts = save_dir.split('/')
parent_id = None
for folder_name in path_parts[4:]:
query = f"mimeType='application/vnd.google-apps.folder' and trashed = false and name='{folder_name}'"
if parent_id:
query += f" and '{parent_id}' in parents"
results = drive_service.files().list(q=query, fields="nextPageToken, files(id, name)").execute()
folder = results.get('files', [])[0]
save_dir_id = folder.get('id')
parent_id = save_dir_id
except HttpError as error:
print(f'An error occurred: {error}')
print(save_dir_id)
return save_dir_id
# /models/Stable-diffusion/sdmodels
def add_shortcut_directory_to_drive (shortcut_name, share_id, save_dir_name, save_dir):
query = f"name='{save_dir_name}' and mimeType='application/vnd.google-apps.folder' and parents in '{save_dir}'"
results = drive_service.files().list(q=query).execute().get('files', [])
if not results:
raise Exception(f"No folder found with name '{save_dir_name}'.")
save_dir_id = results[0]['id']
shortcut_metadata = {
'name': shortcut_name,
'mimeType': 'application/vnd.google-apps.shortcut',
'shortcutDetails': {
'targetId': share_id
},
'parents': [save_dir_id]
}
try:
shortcut_file = drive_service.files().create(body=shortcut_metadata, fields='id').execute()
shortcut_id = shortcut_file['id']
print("Shortcut created successfully!")
except HttpError as error:
print("An error occurred: %s" % error)
shortcut_id = None
return shortcut_id
# /models/Stable-diffusion/*
def add_shortcut_files_to_drive (share_id, save_dir_name, save_dir):
query = f"name = '{save_dir_name}' and mimeType = 'application/vnd.google-apps.folder' and parents in '{save_dir}'"
results = drive_service.files().list(q=query).execute().get('files', [])
if not results:
raise Exception(f"No folder found with name '{save_dir_name}'.")
save_dir_id = results[0]['id']
shortcut_id = None
query = f"'{share_id}' in parents and mimeType != 'application/vnd.google-apps.folder'"
results = drive_service.files().list(q=query).execute().get('files', [])
file_ids = [result['id'] for result in results]
for file_id in file_ids:
file = drive_service.files().get(fileId=file_id).execute()
file_name = file['name']
shortcut_metadata = {
'name': file_name,
'mimeType': 'application/vnd.google-apps.shortcut',
'shortcutDetails': {
'targetId': file_id
},
"parents": [save_dir_id]
}
shortcut_file = drive_service.files().create(body=shortcut_metadata).execute()
shortcut_id = shortcut_file['id']
print(f"Shortcut created for {file_name} with ID: {shortcut_id}")
#delete_shortcut /models/Stable-diffusion/*
def delete_shortcut_files(save_dir):
folder_id = get_save_dir_id(save_dir)
query = f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.shortcut'"
results = drive_service.files().list(q=query, fields="nextPageToken, files(id, name)").execute()
shortcuts = results.get('files', [])
if shortcuts:
print(f"Found {len(shortcuts)} shortcuts in folder {folder_id}:")
for shortcut in shortcuts:
print(f"{shortcut['name']} ({shortcut['id']})")
for shortcut in shortcuts:
drive_service.files().delete(fileId=shortcut['id']).execute()
print(f"Deleted {len(shortcuts)} shortcuts")
else:
print(f"No shortcuts found in folder {folder_id}")