Upload 2 files

src/rep_api.py

@@ -1,213 +1,298 @@
-import replicate 
-import os
-from src.utils import image_to_base64 , update_model_dicts, BB_uploadfile,numpy_to_base64
-from src.deepl import detect_and_translate
-import json
-import time
-style_json="model_dict.json"
-model_dict=json.load(open(style_json,"r"))
-
-
-
-def generate_image_replicate(prompt,lora_model,api_path,aspect_ratio,gallery,model,lora_scale,num_outputs=1,guidance_scale=3.5,seed=None):
-  print(prompt,lora_model,api_path,aspect_ratio)
-
-  #if model=="dev":
-  num_inference_steps=30
-  if model=="schnell":
-    num_inference_steps=5
-
-  if lora_model is not None:
-     api_path=model_dict[lora_model]
-
-  inputs={
-        "model": model,
-        "prompt": detect_and_translate(prompt),
-        "lora_scale":lora_scale,
-        "aspect_ratio": aspect_ratio,
-        "num_outputs":num_outputs,
-        "num_inference_steps":num_inference_steps,
-        "guidance_scale":guidance_scale,
-        "output_format":"png",
-    }
-  if seed is not None:
-    inputs["seed"]=seed
-  output = replicate.run(
-    api_path,
-    input=inputs
-  )
-  print(output)
-  if gallery is None:
-     gallery=[]
-  gallery.append(output[0])
-  return output[0],gallery
-
-
-def replicate_caption_api(image,model,context_text):
-  base64_image = image_to_base64(image)
-  if model=="blip":
-    output = replicate.run(
-    "andreasjansson/blip-2:f677695e5e89f8b236e52ecd1d3f01beb44c34606419bcc19345e046d8f786f9",
-    input={
-        "image": base64_image,
-        "caption": True,
-        "question": context_text,
-        "temperature": 1,
-        "use_nucleus_sampling": False
-    }
-    )
-    print(output)
-
-  elif model=="llava-16":
-    output = replicate.run(
-       # "yorickvp/llava-13b:80537f9eead1a5bfa72d5ac6ea6414379be41d4d4f6679fd776e9535d1eb58bb",
-         "yorickvp/llava-v1.6-34b:41ecfbfb261e6c1adf3ad896c9066ca98346996d7c4045c5bc944a79d430f174",
-        input={
-            "image": base64_image,
-            "top_p": 1,
-            "prompt": context_text,
-            "max_tokens": 1024,
-            "temperature": 0.2
-          }
-        )
-    print(output)
-    output = "".join(output)
-
-  elif model=="img2prompt":
-    output = replicate.run(
-    "methexis-inc/img2prompt:50adaf2d3ad20a6f911a8a9e3ccf777b263b8596fbd2c8fc26e8888f8a0edbb5",
-    input={
-        "image":base64_image
-        }
-    )
-    print(output)
-  return output
-
-def update_replicate_api_key(api_key):
-    os.environ["REPLICATE_API_TOKEN"] = api_key
-    return f"Replicate API key updated: {api_key[:5]}..." if api_key else "Replicate API key cleared"
-
-
-def virtual_try_on(crop, seed, steps, category,  garm_img, human_img,  garment_des):
-    output = replicate.run(
-        "cuuupid/idm-vton:906425dbca90663ff5427624839572cc56ea7d380343d13e2a4c4b09d3f0c30f",
-        input={
-            "crop": crop,
-            "seed": seed,
-            "steps": steps,
-            "category": category,
-           # "force_dc": force_dc,
-            "garm_img": numpy_to_base64( garm_img),
-            "human_img":  numpy_to_base64(human_img),
-            #"mask_only": mask_only,
-            "garment_des": garment_des
-        }
-    )
-    print(output)
-    return output
-
-
-from src.utils import create_zip
-from PIL import Image
-
-
-def process_images(files,model,context_text,token_string):
-    images = []
-    textbox =""
-    for file in files:
-        print(file)
-        image = Image.open(file)
-        if model=="None":
-           caption="[Insert cap here]"
-        else:
-          caption = replicate_caption_api(image,model,context_text)
-        textbox += f"Tags: {caption}, file: " + os.path.basename(file) + "\n"
-        images.append(image)
-        #texts.append(textbox)
-    zip_path=create_zip(files,textbox,token_string)
-    
-    return images, textbox,zip_path
-
-def replicate_create_model(owner,name,visibility="private",hardware="gpu-a40-large"):
-  try:
-    model = replicate.models.create(
-        owner=owner,
-        name=name,
-        visibility=visibility,
-        hardware=hardware,
-    )
-    print(model)
-    return True
-  except Exception as e:
-    print(e)
-    if "A model with that name and owner already exists" in str(e):
-      return True
-    return False
-
-
-
-def traning_function(zip_path,training_model,training_destination,seed,token_string,max_train_steps,hf_repo_id=None,hf_token=None):
-  ##Place holder for now 
-  BB_bucket_name="jarvisdataset" 
-  BB_defult="https://f005.backblazeb2.com/file/"
-  if BB_defult not in zip_path:
-    zip_path=BB_uploadfile(zip_path,os.path.basename(zip_path),BB_bucket_name)
-  print(zip_path)
-  training_logs = f"Using zip traning file at: {zip_path}\n"
-  yield training_logs, None
-  input={
-              "steps": max_train_steps,
-              "lora_rank": 16,
-              "batch_size": 1,
-              "autocaption": True,
-              "trigger_word": token_string,
-              "learning_rate": 0.0004,
-              "seed": seed,
-              "input_images": zip_path
-          }
-  print(training_destination)
-  username,model_name=training_destination.split("/")
-  assert replicate_create_model(username,model_name,visibility="private",hardware="gpu-a40-large"),"Error in creating model on replicate, check API key and username is correct "
-
-  print(input)
-  try:
-      training = replicate.trainings.create(
-          destination=training_destination,
-          version="ostris/flux-dev-lora-trainer:1296f0ab2d695af5a1b5eeee6e8ec043145bef33f1675ce1a2cdb0f81ec43f02",
-          input=input,
-      )
-
-      training_logs = f"Training started with model: {training_model}\n"
-      training_logs += f"Destination: {training_destination}\n"
-      training_logs += f"Seed: {seed}\n"
-      training_logs += f"Token string: {token_string}\n"
-      training_logs += f"Max train steps: {max_train_steps}\n"
-
-      # Poll the training status
-      while training.status != "succeeded":
-          training.reload()
-          training_logs += f"Training status: {training.status}\n"
-          training_logs += f"{training.logs}\n"
-          if training.status == "failed":
-              training_logs += "Training failed!\n"
-              return training_logs, training
-              
-          yield training_logs, None
-          time.sleep(10)  # Wait for 10 seconds before checking again
-        
-      training_logs += "Training completed!\n"
-      if hf_repo_id and hf_token:
-          training_logs += f"Uploading to Hugging Face repo: {hf_repo_id}\n"
-          # Here you would implement the logic to upload to Hugging Face
-
-      traning_finnal=training.output
-
-      # In a real scenario, you might want to download and display some result images
-      # For now, we'll just return the original images
-      #images = [Image.open(file) for file in files]
-      _= update_model_dicts(traning_finnal["version"],token_string,style_json="model_dict.json")
-      traning_finnal["replicate_link"]="https://replicate.com/"+traning_finnal["version"].replace(":","/")
-      yield training_logs, traning_finnal
-
-  except Exception as e:
-        yield f"An error occurred: {str(e)}", None
+import replicate 
+import os
+from src.utils import image_to_base64 , update_model_dicts, BB_uploadfile,numpy_to_base64
+from src.deepl import detect_and_translate
+import json
+import time
+style_json="model_dict.json"
+model_dict=json.load(open(style_json,"r"))
+
+
+def generate_image_control_net(prompt,lora_model,api_path,aspect_ratio,lora_scale,
+                              use_control_net,control_net_type,control_net_img,control_net_strength, 
+                             num_outputs=1,guidance_scale=3.5,seed=None,                
+): 
+   print(prompt,lora_model,api_path,aspect_ratio,use_control_net)
+   inputs = {
+       "prompt": detect_and_translate(prompt),
+        "output_format": "png",
+         "num_outputs":num_outputs,
+        "guidance_scale": guidance_scale,
+        "output_quality": 100,
+   }
+   if seed is not None:
+      inputs["seed"]=seed
+
+
+   if use_control_net:
+      api_path= "xlabs-ai/flux-dev-controlnet:f2c31c31d81278a91b2447a304dae654c64a5d5a70340fba811bb1cbd41019a2" #X labs control net replicate repo
+      lora_url=model_dict[lora_model][1]
+      assert control_net_img is not None, "Please add control net image"
+      control_net_img=image_to_base64(control_net_img)
+      inputs["lora_url"]=lora_url
+      inputs["prompt"]+=", "+model_dict[lora_model][2]
+      inputs["control_image"]=control_net_img
+      inputs["control_type"]=control_net_type
+      inputs["control_strengh"]=control_net_strength
+      #other settings 
+      inputs["steps"]=30
+      inputs["lora_strength"]=lora_scale
+      inputs["negative_prompt"]= "low quality, ugly, distorted, artefacts"
+   else:
+      api_path=model_dict[lora_model][0]
+      inputs["aspect_ratio"]=aspect_ratio
+      inputs["prompt"]+=", "+model_dict[lora_model][2]
+      inputs["num_inference_steps"]=28
+      inputs["model"]="dev"
+      inputs["lora_scale"]=lora_scale
+   
+   #run model
+   output = replicate.run(
+     api_path,
+     input=inputs
+   )
+   print(output)
+   return output[0]
+
+   
+
+
+def generate_image_replicate(prompt,lora_model,api_path,aspect_ratio,model,lora_scale,
+                             num_outputs=1,guidance_scale=3.5,seed=None,
+
+                             ):
+  print(prompt,lora_model,api_path,aspect_ratio)
+
+  #if model=="dev":
+  num_inference_steps=30
+  if model=="schnell":
+    num_inference_steps=5
+
+  if lora_model is not None:
+     api_path=model_dict[lora_model][0]
+
+  inputs={
+        "model": model,
+        "prompt": detect_and_translate(prompt),
+        "lora_scale":lora_scale,
+        "aspect_ratio": aspect_ratio,
+        "num_outputs":num_outputs,
+        "num_inference_steps":num_inference_steps,
+        "guidance_scale":guidance_scale,
+        "output_format":"png",
+    }
+  if seed is not None:
+    inputs["seed"]=seed
+  output = replicate.run(
+    api_path,
+    input=inputs
+  )
+  print(output)
+  return output[0]
+
+
+def replicate_caption_api(image,model,context_text):
+  base64_image = image_to_base64(image)
+  if model=="blip":
+    output = replicate.run(
+    "andreasjansson/blip-2:f677695e5e89f8b236e52ecd1d3f01beb44c34606419bcc19345e046d8f786f9",
+    input={
+        "image": base64_image,
+        "caption": True,
+        "question": context_text,
+        "temperature": 1,
+        "use_nucleus_sampling": False
+    }
+    )
+    print(output)
+
+  elif model=="llava-16":
+    output = replicate.run(
+       # "yorickvp/llava-13b:80537f9eead1a5bfa72d5ac6ea6414379be41d4d4f6679fd776e9535d1eb58bb",
+         "yorickvp/llava-v1.6-34b:41ecfbfb261e6c1adf3ad896c9066ca98346996d7c4045c5bc944a79d430f174",
+        input={
+            "image": base64_image,
+            "top_p": 1,
+            "prompt": context_text,
+            "max_tokens": 1024,
+            "temperature": 0.2
+          }
+        )
+    print(output)
+    output = "".join(output)
+
+  elif model=="img2prompt":
+    output = replicate.run(
+    "methexis-inc/img2prompt:50adaf2d3ad20a6f911a8a9e3ccf777b263b8596fbd2c8fc26e8888f8a0edbb5",
+    input={
+        "image":base64_image
+        }
+    )
+    print(output)
+  return output
+
+def update_replicate_api_key(api_key):
+    os.environ["REPLICATE_API_TOKEN"] = api_key
+    return f"Replicate API key updated: {api_key[:5]}..." if api_key else "Replicate API key cleared"
+
+
+def virtual_try_on(crop, seed, steps, category,  garm_img, human_img,  garment_des):
+    output = replicate.run(
+        "cuuupid/idm-vton:906425dbca90663ff5427624839572cc56ea7d380343d13e2a4c4b09d3f0c30f",
+        input={
+            "crop": crop,
+            "seed": seed,
+            "steps": steps,
+            "category": category,
+           # "force_dc": force_dc,
+            "garm_img": numpy_to_base64( garm_img),
+            "human_img":  numpy_to_base64(human_img),
+            #"mask_only": mask_only,
+            "garment_des": garment_des
+        }
+    )
+    print(output)
+    return output
+
+
+from src.utils import create_zip
+from PIL import Image
+
+
+def process_images(files,model,context_text,token_string):
+    images = []
+    textbox =""
+    for file in files:
+        print(file)
+        image = Image.open(file)
+        if model=="None":
+           caption="[Insert cap here]"
+        else:
+          caption = replicate_caption_api(image,model,context_text)
+        textbox += f"Tags: {caption}, file: " + os.path.basename(file) + "\n"
+        images.append(image)
+        #texts.append(textbox)
+    zip_path=create_zip(files,textbox,token_string)
+    
+    return images, textbox,zip_path
+
+def replicate_create_model(owner,name,visibility="private",hardware="gpu-a40-large"):
+  try:
+    model = replicate.models.create(
+        owner=owner,
+        name=name,
+        visibility=visibility,
+        hardware=hardware,
+    )
+    print(model)
+    return True
+  except Exception as e:
+    print(e)
+    if "A model with that name and owner already exists" in str(e):
+      return True
+    return False
+
+
+
+def traning_function(zip_path,training_model,training_destination,seed,token_string,max_train_steps,hf_repo_id=None,hf_token=None):
+  ##Place holder for now 
+  BB_bucket_name="jarvisdataset" 
+  BB_defult="https://f005.backblazeb2.com/file/"
+  if BB_defult not in zip_path:
+    zip_path=BB_uploadfile(zip_path,os.path.basename(zip_path),BB_bucket_name)
+  print(zip_path)
+  training_logs = f"Using zip traning file at: {zip_path}\n"
+  yield training_logs, None
+  input={
+              "steps": max_train_steps,
+              "lora_rank": 16,
+              "batch_size": 1,
+              "autocaption": True,
+              "trigger_word": token_string,
+              "learning_rate": 0.0004,
+              "seed": seed,
+              "input_images": zip_path
+          }
+  print(training_destination)
+  username,model_name=training_destination.split("/")
+  assert replicate_create_model(username,model_name,visibility="private",hardware="gpu-a40-large"),"Error in creating model on replicate, check API key and username is correct "
+
+  print(input)
+  try:
+      training = replicate.trainings.create(
+          destination=training_destination,
+          version="ostris/flux-dev-lora-trainer:1296f0ab2d695af5a1b5eeee6e8ec043145bef33f1675ce1a2cdb0f81ec43f02",
+          input=input,
+      )
+
+      training_logs = f"Training started with model: {training_model}\n"
+      training_logs += f"Destination: {training_destination}\n"
+      training_logs += f"Seed: {seed}\n"
+      training_logs += f"Token string: {token_string}\n"
+      training_logs += f"Max train steps: {max_train_steps}\n"
+
+      # Poll the training status
+      while training.status != "succeeded":
+          training.reload()
+          training_logs += f"Training status: {training.status}\n"
+          training_logs += f"{training.logs}\n"
+          if training.status == "failed":
+              training_logs += "Training failed!\n"
+              return training_logs, training
+              
+          yield training_logs, None
+          time.sleep(10)  # Wait for 10 seconds before checking again
+        
+      training_logs += "Training completed!\n"
+      if hf_repo_id and hf_token:
+          training_logs += f"Uploading to Hugging Face repo: {hf_repo_id}\n"
+          # Here you would implement the logic to upload to Hugging Face
+
+      traning_finnal=training.output
+
+      # In a real scenario, you might want to download and display some result images
+      # For now, we'll just return the original images
+      #images = [Image.open(file) for file in files]
+      _= update_model_dicts(traning_finnal["version"],token_string,style_json="model_dict.json")
+      traning_finnal["replicate_link"]="https://replicate.com/"+traning_finnal["version"].replace(":","/")
+      yield training_logs, traning_finnal
+
+  except Exception as e:
+        yield f"An error occurred: {str(e)}", None
+
+
+def sam_segment(image,prompt,negative_prompt,adjustment_factor=-15):
+  #img2 base64
+  image = image_to_base64(image)
+  output = replicate.run(
+      "schananas/grounded_sam:ee871c19efb1941f55f66a3d7d960428c8a5afcb77449547fe8e5a3ab9ebc21c",
+      input={
+          "image": image,
+          "mask_prompt": prompt,
+          "adjustment_factor": adjustment_factor,
+          "negative_mask_prompt":negative_prompt
+      }
+  )
+  out_items={}
+  for item in output:
+    # https://replicate.com/schananas/grounded_sam/api#output-schema
+    print(item)
+    out_items[os.path.basename(item).split(".")[0]]=item
+  return out_items 
+
+
+def replicate_zest(img,material_img="https://replicate.delivery/pbxt/Kl23gJODaW7EuxrDzBG9dcgqRdMaYSWmBQ9UexnwPiL7AnIr/3.jpg"):
+  if type(img)!=str:
+    img=image_to_base64(img)
+  if type(material_img)!=str:
+    material_img=image_to_base64(material_img)
+
+  output = replicate.run(
+      "camenduru/zest:11abc0a411459327938957581151c642dd1bee4cefe443a9a63b230c4fbc0952",
+      input={
+          "input_image": img,
+          "material_image":material_img
+      }
+  ) 
+  print(output)
+  return output

src/utils.py

@@ -1,100 +1,138 @@
-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"
-
+import b2sdk.v2 as b2 #Backblaze img2img upload bucket
+import base64
+import os
+from io import BytesIO
+import requests
+import io
+from PIL import Image
+import numpy as npzipfile
+import zipfile
+import numpy as np
+import json 
+import gradio as gr
+
+
+def convert_to_pil(img):
+    if isinstance(img, np.ndarray):
+        img = Image.fromarray(img)
+    # If the image is a URL, fetch the image and convert it to a PIL image
+    elif isinstance(img, str) and img.startswith('http'):
+        response = requests.get(img)
+        img = Image.open(BytesIO(response.content))
+    return img 
+
+def open_image_from_url(image_url):
+  response = requests.get(image_url)
+  img = Image.open(BytesIO(response.content))
+  return img
+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 cut_alpha_from_image(image):
+  """Cuts out the alpha channel from a Pillow image, returning a new image.
+
+  Args:
+    image: The Pillow image to cut the alpha channel from.
+
+  Returns:
+    A new Pillow image with the alpha channel removed.
+  """
+  if image.mode == 'RGBA':
+    # Extract RGB channels and create a new image without alpha
+    rgb_image = Image.new("RGB", image.size, (0, 0, 0))
+    rgb_image.paste(image, mask=image.split()[3])  # Use alpha channel as mask
+    return rgb_image
+  else:
+    return image  # Image doesn't have an alpha channel, return as is
+
+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()
+  if isinstance(img, np.ndarray):
+     img=Image.fromarray(img)
+  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"
+