Upload script.py

script.py

@@ -0,0 +1,129 @@
+import sys
+import subprocess
+from safetensors.torch import load_file
+from diffusers import AutoPipelineForText2Image
+from datasets import load_dataset
+from huggingface_hub.repocard import RepoCard
+from huggingface_hub import HfApi
+import torch
+import re
+import argparse
+import os
+import zipfile
+
+def do_preprocess(class_data_dir):
+    print("Unzipping dataset")
+    zip_file_path = f"{class_data_dir}/class_images.zip"
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        zip_ref.extractall(class_data_dir)
+    os.remove(zip_file_path)
+
+def do_train(script_args):
+    # Pass all arguments to trainer.py
+    print("Starting training...")
+    result = subprocess.run(['python', 'trainer.py'] + script_args)
+    if result.returncode != 0:
+        raise Exception("Training failed.")
+
+def replace_output_dir(text, output_dir, replacement):
+            # Define a pattern that matches the output_dir followed by whitespace, '/', new line, or "'"
+            # Add system name from HF only in the correct spots
+            pattern = rf"{output_dir}(?=[\s/'\n])"
+            return re.sub(pattern, replacement, text)
+    
+def do_inference(dataset_name, output_dir, num_tokens):
+    widget_content = []
+    try:
+        print("Starting inference to generate example images...")
+        dataset = load_dataset(dataset_name)
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            "./lora-ease-wsl/tPonynai3_v6.safetensors", torch_dtype=torch.float16
+        )
+        pipe = pipe.to("cuda")
+        pipe.load_lora_weights(f'{output_dir}/pytorch_lora_weights.safetensors')
+        
+        prompts = dataset["train"]["prompt"]
+        if(num_tokens > 0):
+            tokens_sequence = ''.join(f'<s{i}>' for i in range(num_tokens))
+            tokens_list = [f'<s{i}>' for i in range(num_tokens)]
+        
+            state_dict = load_file(f"{output_dir}/{output_dir}_emb.safetensors")
+            pipe.load_textual_inversion(state_dict["clip_l"], token=tokens_list, text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+            pipe.load_textual_inversion(state_dict["clip_g"], token=tokens_list, text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+            
+            prompts = [prompt.replace("TOK", tokens_sequence) for prompt in prompts]    
+
+        for i, prompt in enumerate(prompts):
+            image = pipe(prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
+            filename = f"image-{i}.png"
+            image.save(f"{output_dir}/{filename}")
+            card_dict = {
+                "text": prompt,
+                "output": {
+                    "url": filename
+                }
+            }
+            widget_content.append(card_dict)
+    except Exception as e:
+        print("Something went wrong with generating images, specifically: ", e)
+    
+    try:
+        api = HfApi()
+        username = api.whoami()["name"]
+        repo_id = api.create_repo(f"{username}/{output_dir}", exist_ok=True, private=True).repo_id
+        
+        with open(f'{output_dir}/README.md', 'r') as file:
+            readme_content = file.read()
+        
+    
+        readme_content = replace_output_dir(readme_content, output_dir, f"{username}/{output_dir}")
+        
+        card = RepoCard(readme_content)
+        if widget_content: 
+            card.data["widget"] = widget_content
+            card.save(f'{output_dir}/README.md')
+    
+        print("Starting upload...")
+        api.upload_folder(
+            folder_path=output_dir,
+            repo_id=f"{username}/{output_dir}",
+            repo_type="model",
+        )
+    except Exception as e:
+        print("Something went wrong with uploading your model, specificaly: ", e)
+    else:
+        print("Upload finished!")
+
+import sys
+import argparse
+
+def main():
+    # Capture all arguments except the script name
+    script_args = sys.argv[1:]
+
+    # Create the argument parser
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset_name', required=True)
+    parser.add_argument('--output_dir', required=True)
+    parser.add_argument('--num_new_tokens_per_abstraction', type=int, default=0)
+    parser.add_argument('--train_text_encoder_ti', action='store_true')
+    parser.add_argument('--class_data_dir', help="Name of the class images dataset")
+
+    # Parse known arguments
+    args, _ = parser.parse_known_args(script_args)
+
+    # Set num_tokens to 0 if '--train_text_encoder_ti' is not present
+    if not args.train_text_encoder_ti:
+        args.num_new_tokens_per_abstraction = 0
+
+    # Proceed with training and inference
+    if args.class_data_dir:
+        do_preprocess(args.class_data_dir)
+        print("Pre-processing finished!")
+    do_train(script_args)
+    print("Training finished!")
+    do_inference(args.dataset_name, args.output_dir, args.num_new_tokens_per_abstraction)
+    print("All finished!")
+
+if __name__ == "__main__":
+    main()