Upload 2 files

src/rep_api.py

@@ -1,13 +1,11 @@
 import replicate 
 import os
-from src.utils import image_to_base64
-from src.utils import BB_uploadfile
+from src.utils import image_to_base64 , update_model_dicts, BB_uploadfile
 import json
 import time
 style_json="model_dict.json"
 model_dict=json.load(open(style_json,"r"))
-from dotenv import load_dotenv, find_dotenv
-_ = load_dotenv(find_dotenv())
+
 
 
 def generate_image_replicate(prompt,lora_model,api_path,aspect_ratio,gallery,model="dev",lora_scale=1,num_outputs=1,guidance_scale=3.5,seed=None):
@@ -19,8 +17,7 @@ def generate_image_replicate(prompt,lora_model,api_path,aspect_ratio,gallery,mod
 
   if lora_model is not None:
      api_path=model_dict[lora_model]
-  
-  
+
   inputs={
         "model": model,
         "prompt": prompt,
@@ -91,26 +88,46 @@ from src.utils import create_zip
 from PIL import Image
 
 
-def process_images(files,model,context_text):
+def process_images(files,model,context_text,token_string):
     images = []
     textbox =""
     for file in files:
         print(file)
         image = Image.open(file)
-        caption = replicate_caption_api(image,model,context_text)
+        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,"TOK")
+    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"
@@ -126,6 +143,9 @@ def traning_function(zip_path,training_model,training_destination,seed,token_str
               "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(
@@ -162,6 +182,8 @@ def traning_function(zip_path,training_model,training_destination,seed,token_str
       # 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:

src/utils.py

@@ -36,7 +36,7 @@ def add_to_prompt(existing_prompt, new_prompt):
 def image_to_base64(img):
   buffered = io.BytesIO()
   img.save(buffered, format="PNG")
-  img_str = base64.b64encode(buffered.getvalue()).decode('utexf-8')
+  img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
   return "data:image/png;base64,"+img_str
 
 def create_zip(files,captions,trigger):
@@ -44,26 +44,33 @@ def create_zip(files,captions,trigger):
     captions=captions.split("\n")
     #cute files and "tags:"
     captions= [cap.split("file:")[0][5:] for cap in captions]
-    with tempfile.TemporaryDirectory() as temp_dir:
-      # Create a zip file
-      zip_path = os.path.join(temp_dir, "training_data.zip")
-      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"
-            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
-
+    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,application_key_id = os.getenv("BB_KeyID"),application_key = os.getenv("BB_AppKey"),):
+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/"