Upload 3 files

app.py

@@ -1,238 +1,145 @@
-import gradio as gr
-from PIL import Image
-import os
-import replicate
-import time 
-import base64
-import numpy as npzipfile
-import tempfile
-import zipfile
-import b2sdk.v2 as b2 #Backblaze img2img upload bucket
-import shutil
-import requests
-import io
-
-
-
-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 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")
-    print(zip_path)
-    return images, textbox
-
-def BB_uploadfile(b2_api,local_file,file_name,BB_bucket_name,FRIENDLY_URL=True):
-    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
-
-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]
-    #temp_dir="/content"
-    #os.makedirs(temp_dir, exist_ok=True)
-    # Create a zip file
-    #os.makedirs(temp_dir, exist_ok=True)
-    zip_path = "training_data.zip" #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)
-
-    file_url= BB_uploadfile(b2_api,zip_path,f"training_data_{trigger}.zip",BB_bucket_name)
-    return file_url
-
-
-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 traning_function(files,text_output,training_model,training_destination,seed,token_string,max_train_steps,hf_repo_id,hf_token):
-  print(files,text_output)
-  zip_path = create_zip(files,text_output,token_stringn)
-  print(zip_path)
-  training_logs = f"Created zip file at: {zip_path}\n"
-  yield training_logs, None
-  
-  try:
-      training = replicate.trainings.create(
-          destination=training_destination,
-          version="ostris/flux-dev-lora-trainer:1296f0ab2d695af5a1b5eeee6e8ec043145bef33f1675ce1a2cdb0f81ec43f02",
-          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
-          },
-      )
-
-      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"
-          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
-
-      # 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]
-      
-      yield training_logs, images
-
-  except Exception as e:
-        yield f"An error occurred: {str(e)}", None
-  
-with gr.Blocks() as demo:
-    gr.Markdown("# Image Captioning")
-    with gr.Row():
-        input_images = gr.File(file_count="multiple", type="filepath", label="Upload Images")
-        label_model = gr.Dropdown(["blip", "llava-16","img2prompt"], label="Caption model", info="Auto caption model")
-        context_text = gr.Textbox(label="Context Text", info="Context Text for auto catpion",value=" I want a description caption for this image")
-        # Replicate API Key input
-        replicate_api_key = gr.Textbox(
-            label="Replicate API Key", 
-            info="API key for Replicate", 
-            type="password"
-        )
-        api_key_status = gr.Textbox(label="API Key Status", interactive=False)
-
-    with gr.Row():
-        process_button = gr.Button("Process Images")
-    #Image outputs 
-    with gr.Row():
-      gr.Markdown("# Captions")
-    with gr.Row():
-      with gr.Column():
-        image_output = gr.Gallery(type="pil",object_fit="fill")
-      with gr.Column():
-        text_output = gr.Textbox( interactive=True)
-    #Traning options 
-    with gr.Row():
-      gr.Markdown("# Training on replicate")
-    with gr.Row():
-      traning_model = gr.Dropdown(["flux", "SDXL",""], label="Caption model", info="Auto caption model")
-      traning_destination = gr.Textbox(label="destination",info="add in replicate model destination")
-      seed = gr.Number(label="Seed", value=42,info="Random seed integer for reproducible training. Leave empty to use a random seed.")
-      token_stringn = "TOK"# gr.Textbox(label="Token string",value="TOK",info="A unique string that will be trained to refer to the concept in the input images. Can be anything, but TOK works well.")
-      max_train_steps =gr.Number(label="max_train_steps", value= 1000, info="Number of individual training steps. Takes precedence over num_train_epochs.")
-    with gr.Row():
-      hf_repo_id = gr.Textbox(label="Hugging face repo id",info="Hugging Face repository ID, if you'd like to upload the trained LoRA to Hugging Face. For example, lucataco/flux-dev-lora.")
-      hf_token = gr.Textbox(label="Hugging face write token",info="Hugging Face token, if you'd like to upload the trained LoRA to Hugging Face.")
-    with gr.Row():
-      train_button = gr.Button("Train")
-    with gr.Row():
-      training_logs = gr.Textbox(label="Training logs")
-      training_images = gr.Gallery(label="Training images")
-
-
-
-    train_button.click(fn=traning_function, inputs=[input_images,text_output,traning_model,traning_destination,seed,token_stringn,max_train_steps,hf_repo_id,hf_token], 
-                       outputs=[image_output,text_output])
-    
-    process_button.click(fn=process_images, inputs=[input_images,label_model,context_text], outputs=[image_output,text_output])
-    # Add event listener for API key changes
-    replicate_api_key.change(
-        fn=update_replicate_api_key,
-        inputs=[replicate_api_key],
-        outputs=[api_key_status]
-    )
-
-
+
+from dotenv import load_dotenv, find_dotenv
+_ = load_dotenv(find_dotenv())
+from src.utils import create_zip,add_to_prompt,update_dropdown
+from src.rep_api import replicate_caption_api,generate_image_replicate,traning_function,update_replicate_api_key
+import gradio as gr
+from PIL import Image
+import os
+import time
+import json
+
+
+# The dictionary data
+prompt_dict = {
+    "Character": ["Asian girl with black hair", "A man with blond hair", "A Cat girl anime character with purple hair", "A Green Alien with big black eyes"],
+    "Clothes": ["Wearing a blue jacket", "Wearing a black business suit", "Wearing a purple jumpsuit", "Wearing shorts and a white T-shirt"],
+    "Pose": ["Close up portrait", "Standing doing a peace sign", "Folding arms", "holding a phone"],
+    "Style": ["Simple white background", "Fashion runway", "Inside a business conference", "Inside a spaceship"]
+}
+style_json="model_dict.json"
+model_dict=json.load(open(style_json,"r"))
+
+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
+
+
+
+
+with gr.Blocks( theme="NoCrypt/miku") as demo:
+   
+    with gr.Tabs() as tabs:
+      with gr.TabItem("Image Generator"):
+        gr.Markdown(" #Image Generator")
+        with gr.Row():
+          with gr.Column():
+            inp = gr.Textbox(label="Prompt")
+
+
+            btn = gr.Button("Generate")
+          with gr.Column():
+            ar = gr.Dropdown(["1:1","16:9","9:16","5:3"], label="Aspect Ratio", info="Aspect Ratio")
+            style_mode = gr.Dropdown(model_dict.keys(),label="Style lore")
+            api_path = gr.Textbox(label="API_route",info="replicate api route goes here")
+
+
+
+        with gr.Accordion("Prompt Support", open=False):
+          for key, values in prompt_dict.items():
+                with gr.Row():
+                    #gr.Markdown(f"**{key}**")
+                    gr.Button(key,interactive=False)
+                    for value in values:
+                        gr.Button(value).click(add_to_prompt, inputs=[inp, gr.Textbox(value,visible=False)], outputs=inp)
+
+        with gr.Row():
+          gen_out = gr.Image(label="Generated Image",type="filepath")
+
+
+        btn.click(generate_image_replicate, inputs=[inp,api_path], outputs=gen_out,queue=True)
+
+
+
+      with gr.TabItem("Model Trainner"):
+        gr.Markdown("# Image Importing & Auto captions")
+        with gr.Row():
+            input_images = gr.File(file_count="multiple", type="filepath", label="Upload Images")
+            label_model = gr.Dropdown(["blip", "llava-16","img2prompt"], label="Caption model", info="Auto caption model")
+            token_string= gr.Textbox(label="Token string",value="TOK",interactive=True,
+                                     info="A unique string that will be trained to refer to the concept in the input images. Can be anything, but TOK works well.")
+            context_text = gr.Textbox(label="Context Text", info="Context Text for auto catpion",value=" I want a description caption for this image")
+            # Replicate API Key input
+            replicate_api_key = gr.Textbox(
+                label="Replicate API Key",
+                info="API key for Replicate",
+                value=os.environ.get("REPLICATE_API_TOKEN", ""),
+                type="password"
+            )
+            api_key_status = gr.Textbox(label="API Key Status", interactive=False)
+
+        with gr.Row():
+            process_button = gr.Button("Process Images")
+        #Image outputs
+        with gr.Row():
+          gr.Markdown("# Traning Captions Data")
+        with gr.Row():
+          with gr.Column():
+            image_output = gr.Gallery(type="pil",object_fit="fill")
+          with gr.Column():
+            text_output = gr.Textbox( interactive=True)
+        with gr.Row():
+            zip_output = gr.File(label="Zip file")
+            btn_update_zip = gr.Button("Update zip file")
+
+
+
+        #Traning options
+        with gr.Row():
+          gr.Markdown("# Training on replicate")
+        with gr.Row():
+          traning_model = gr.Dropdown(["flux"], label="Caption model", info="Auto caption model")
+          traning_destination = gr.Textbox(label="destination",info="add in replicate model destination")
+          seed = gr.Number(label="Seed", value=42,info="Random seed integer for reproducible training. Leave empty to use a random seed.")
+
+          max_train_steps =gr.Number(label="max_train_steps", value= 1000, info="Number of individual training steps. Takes precedence over num_train_epochs.")
+        #with gr.Row():
+        # hf_repo_id = gr.Textbox(label="Hugging face repo id",info="Hugging Face repository ID, if you'd like to upload the trained LoRA to Hugging Face. For example, lucataco/flux-dev-lora.")
+        # hf_token = gr.Textbox(label="Hugging face write token",info="Hugging Face token, if you'd like to upload the trained LoRA to Hugging Face.")
+        with gr.Row():
+          train_button = gr.Button("Train")
+        with gr.Row():
+          training_logs = gr.Textbox(label="Training logs")
+          traning_finnal = gr.Textbox(label="Traning finnal")
+          #training_images = gr.Gallery(label="Training images")
+
+
+        #gr.Textbox("TOK",visible=False) added to deal with odd ies of the token string being a gradio class
+        train_button.click(fn=traning_function, inputs=[zip_output,traning_model,traning_destination,seed,token_string,max_train_steps], #,hf_repo_id,hf_token
+                          outputs=[training_logs,traning_finnal],queue=True)
+        process_button.click(fn=process_images, inputs=[input_images,label_model,context_text,token_string], outputs=[image_output,text_output,zip_output],queue=True)
+        btn_update_zip.click(fn=create_zip, inputs=[image_output,text_output,token_string],outputs=zip_output)
+        # Add event listener for API key changes
+        traning_finnal.change(
+            fn=update_dropdown,
+            inputs=[traning_finnal,token_string],
+            outputs=style_mode
+        )
+        replicate_api_key.change(
+            fn=update_replicate_api_key,
+            inputs=[replicate_api_key],
+            outputs=[api_key_status]
+        )
+#jarvis-labs2024/sioux-flux
+demo.queue() # Queue for concurrent users 
 demo.launch(debug=True)

model_dict.json

@@ -0,0 +1,7 @@
+{
+    "Base": "black-forest-labs/flux-dev",
+    "Raylean": "jarvis-labs2024/flux-raylene:5574556226d11e0f10855a957d91f118a9178c8fc77e7e7b18830627ce3184f1",
+    "Alice": "jarvis-labs2024/flux-raylene:5574556226d11e0f10855a957d91f118a9178c8fc77e7e7b18830627ce3184f1",
+    "AppleSeed": "jarvis-labs2024/flux-appleseed:0aecb9fdfb17a2517112cc70b4a1898aa7791da84a010419782ce7043481edec",
+    "console_cowboy_flux": "jarvis-labs2024/console_cowboy_flux:53ff894d719f73dc11ca54fdb6ecf044d7d202aa30fce43236fbfda30b19ef62"
+}

requirments.txt

@@ -0,0 +1,6 @@
+gradio
+fal
+fal-client
+numpy
+replicate
+python-dotenv