mirror of
https://github.com/KoboldAI/KoboldAI-Client.git
synced 2025-06-05 21:59:24 +02:00
Better downloading for unsaved or partially saved files
This commit is contained in:
35
aiserver.py
35
aiserver.py
@@ -8354,17 +8354,23 @@ def UI_2_save_story(data):
|
|||||||
#We have an ack that it's OK to save over the file if one exists
|
#We have an ack that it's OK to save over the file if one exists
|
||||||
koboldai_vars.save_story()
|
koboldai_vars.save_story()
|
||||||
|
|
||||||
def directory_to_zip_data(directory: str) -> bytes:
|
def directory_to_zip_data(directory: str, overrides: Optional[dict]) -> bytes:
|
||||||
|
overrides = overrides or {}
|
||||||
buffer = BytesIO()
|
buffer = BytesIO()
|
||||||
|
|
||||||
with zipfile.ZipFile(buffer, "w") as zipf:
|
with zipfile.ZipFile(buffer, "w") as zipf:
|
||||||
for root, _, files in os.walk(directory):
|
for root, _, files in os.walk(directory):
|
||||||
for file in files:
|
for file in files:
|
||||||
p = os.path.join(root, file)
|
p = os.path.join(root, file)
|
||||||
zipf.write(
|
z_path = os.path.join(*p.split(os.path.sep)[2:])
|
||||||
p,
|
|
||||||
os.path.join(*p.split(os.path.sep)[2:])
|
if z_path in overrides:
|
||||||
)
|
continue
|
||||||
|
|
||||||
|
zipf.write(p, z_path)
|
||||||
|
|
||||||
|
for path, contents in overrides.items():
|
||||||
|
zipf.writestr(path, contents)
|
||||||
|
|
||||||
return buffer.getvalue()
|
return buffer.getvalue()
|
||||||
|
|
||||||
@@ -8374,8 +8380,25 @@ def directory_to_zip_data(directory: str) -> bytes:
|
|||||||
@app.route("/story_download")
|
@app.route("/story_download")
|
||||||
@logger.catch
|
@logger.catch
|
||||||
def UI_2_download_story():
|
def UI_2_download_story():
|
||||||
|
save_exists = path.exists(koboldai_vars.save_paths.base)
|
||||||
|
if koboldai_vars.gamesaved and save_exists:
|
||||||
|
# Disk is up to date; download from disk
|
||||||
|
data = directory_to_zip_data(koboldai_vars.save_paths.base)
|
||||||
|
elif save_exists:
|
||||||
|
# We aren't up to date but we are saved; patch what disk gives us
|
||||||
|
data = directory_to_zip_data(
|
||||||
|
koboldai_vars.save_paths.base,
|
||||||
|
{"story.json": koboldai_vars.to_json("story_settings")}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# We are not saved; send json in zip from memory
|
||||||
|
buffer = BytesIO()
|
||||||
|
with zipfile.ZipFile(buffer, "w") as zipf:
|
||||||
|
zipf.writestr("story.json", koboldai_vars.to_json("story_settings"))
|
||||||
|
data = buffer.getvalue()
|
||||||
|
|
||||||
return Response(
|
return Response(
|
||||||
directory_to_zip_data(koboldai_vars.save_paths.base),
|
data,
|
||||||
mimetype="application/octet-stream",
|
mimetype="application/octet-stream",
|
||||||
headers={"Content-disposition": f"attachment; filename={koboldai_vars.story_name}.kaistory"}
|
headers={"Content-disposition": f"attachment; filename={koboldai_vars.story_name}.kaistory"}
|
||||||
)
|
)
|
||||||
|
Reference in New Issue
Block a user