Update src/utils.py

src/utils.py

@@ -1,90 +1,100 @@
-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 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"
-
+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"
+