Upload 2 files

src/rep_api.py

@@ -1,25 +1,32 @@
 import replicate 
 import os
 from src.utils import image_to_base64
-from src.utils import BB_uploadfile, BB_bucket_name, BB_defurl, b2_api
-
+from src.utils import 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())
+
 
-#api_path= "jarvis-labs2024/flux-raylene:5574556226d11e0f10855a957d91f118a9178c8fc77e7e7b18830627ce3184f1"
-#api_path= "jarvis-labs2024/flux-raylene:5574556226d11e0f10855a957d91f118a9178c8fc77e7e7b18830627ce3184f1"
-def generate_image_replicate(prompt,api_path,aspect_ratio="1:1",model="dev",lora_scale=1,num_outputs=1,guidance_scale=3.5,seed=None):
-  print(prompt,api_path)
+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):
+  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": prompt,
         "lora_scale":lora_scale,
         "aspect_ratio": aspect_ratio,
         "num_outputs":num_outputs,
-       
         "num_inference_steps":num_inference_steps,
         "guidance_scale":guidance_scale,
     }
@@ -30,7 +37,10 @@ def generate_image_replicate(prompt,api_path,aspect_ratio="1:1",model="dev",lora
     input=inputs
   )
   print(output)
-  return output[0]
+  if gallery is None:
+     gallery=[]
+  gallery.append(output[0])
+  return output[0],gallery
 
 
 def replicate_caption_api(image,model,context_text):
@@ -77,13 +87,31 @@ 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"
 
+from src.utils import create_zip
+from PIL import Image
 
 
+def process_images(files,model,context_text):
+    images = []
+    textbox =""
+    for file in files:
+        print(file)
+        image = Image.open(file)
+        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")
+    
+    return images, textbox,zip_path
+
 def traning_function(zip_path,training_model,training_destination,seed,token_string,max_train_steps,hf_repo_id=None,hf_token=None):
-  #print(files,text_output)
-  #zip_path = create_zip(files,text_output,token_stringn)
-  if BB_defurl not in zip_path:
-    zip_path=BB_uploadfile(b2_api,zip_path,os.path.basename(zip_path),BB_bucket_name)
+  ##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

src/utils.py

@@ -57,18 +57,17 @@ def create_zip(files,captions,trigger):
             caption_content = captions[i] +f", {trigger}"
             zip_file.writestr(caption_name, caption_content)
     return zip_path
-info = b2.InMemoryAccountInfo()
-b2_api = b2.B2Api(info)
-application_key_id = os.getenv("BB_KeyID")
-application_key = os.getenv("BB_AppKey")
-#print(application_key_id,application_key)
-b2_api.authorize_account("production", application_key_id, application_key)
-BB_bucket_name=os.getenv("BB_bucket")
-BB_bucket=b2_api.get_bucket_by_name(os.getenv("BB_bucket"))
-BB_defurl="https://f005.backblazeb2.com/file/"
 
 
-def BB_uploadfile(b2_api,local_file,file_name,BB_bucket_name,FRIENDLY_URL=True):
+
+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"),):
+    info = b2.InMemoryAccountInfo()
+    b2_api = b2.B2Api(info)
+    #print(application_key_id,application_key)
+    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,