Update src/rep_api.py

src/rep_api.py

@@ -1,190 +1,210 @@
-import replicate 
-import os
-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"))
-
-
-
-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,
-    }
-  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"
-
-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
+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=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,
+    }
+  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