Spaces:
Sleeping
Sleeping
| import b2sdk.v2 as b2 #Backblaze img2img upload bucket | |
| import base64 | |
| import shutil | |
| import os | |
| import requests | |
| import io | |
| from PIL import Image | |
| import numpy as npzipfile | |
| import zipfile | |
| import tempfile | |
| import json | |
| import gradio as gr | |
| def update_model_dicts(traning_finnal,token_string,style_json="model_dict.json"): | |
| print(traning_finnal,token_string) | |
| current_style_dict=json.load(open(style_json,"r")) | |
| current_style_dict[token_string]=traning_finnal | |
| with open(style_json, "w") as json_file: | |
| json.dump(current_style_dict, json_file, indent=4) | |
| json_file.close() | |
| # Return the updated dictionary keys for updating the Dropdown | |
| return list(current_style_dict.keys()) | |
| def update_dropdown(traning_finnal, token_string): | |
| updated_keys = update_model_dicts(traning_finnal, token_string) | |
| return gr.Dropdown.update(choices=updated_keys) | |
| def add_to_prompt(existing_prompt, new_prompt): | |
| if existing_prompt: | |
| return f"{existing_prompt}, {new_prompt}" | |
| else: | |
| return new_prompt | |
| def numpy_to_base64(image_np): | |
| """Converts a numpy image to base64 string.""" | |
| img = Image.fromarray(image_np) | |
| buffered = io.BytesIO() | |
| img.save(buffered, format="PNG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode('utf-8') | |
| return "data:image/png;base64,"+img_str | |
| def image_to_base64(img): | |
| buffered = io.BytesIO() | |
| img.save(buffered, format="PNG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode('utf-8') | |
| return "data:image/png;base64,"+img_str | |
| def create_zip(files,captions,trigger): | |
| #Caption processing | |
| captions=captions.split("\n") | |
| #cute files and "tags:" | |
| captions= [cap.split("file:")[0][5:] for cap in captions] | |
| print("files",len(files),"captions",len(captions)) | |
| #assert len(files)==len(captions) , "File amount does not equal the captions amount please check" | |
| temp_dir="./datasets/" | |
| os.makedirs(temp_dir,exist_ok=True) | |
| zip_path = os.path.join(temp_dir, f"training_data_{trigger}.zip") | |
| if os.path.exists(zip_path): | |
| os.remove(zip_path) | |
| with zipfile.ZipFile(zip_path, "w") as zip_file: | |
| for i, file in enumerate(files): | |
| # Add image to zip | |
| image_name = f"image_{i}.jpg" | |
| print(file) | |
| zip_file.write(file, image_name) | |
| # Add caption to zip | |
| caption_name = f"image_{i}.txt" | |
| caption_content = captions[i] +f", {trigger}" | |
| zip_file.writestr(caption_name, caption_content) | |
| return zip_path | |
| def BB_uploadfile(local_file,file_name,BB_bucket_name,FRIENDLY_URL=True): | |
| info = b2.InMemoryAccountInfo() | |
| b2_api = b2.B2Api(info) | |
| #print(application_key_id,application_key) | |
| application_key_id = os.getenv("BB_KeyID") | |
| application_key = os.getenv("BB_AppKey") | |
| b2_api.authorize_account("production", application_key_id, application_key) | |
| BB_bucket=b2_api.get_bucket_by_name(BB_bucket_name) | |
| BB_defurl="https://f005.backblazeb2.com/file/" | |
| metadata = {"key": "value"} | |
| uploaded_file = BB_bucket.upload_local_file( | |
| local_file=local_file, | |
| file_name=file_name, | |
| file_infos=metadata, | |
| ) | |
| img_url=b2_api.get_download_url_for_fileid(uploaded_file.id_) | |
| if FRIENDLY_URL: #Get friendly URP | |
| img_url=BB_defurl+BB_bucket_name+"/"+file_name | |
| print("backblaze", img_url) | |
| return img_url | |
| #file="/content/training_data.zip" | |