elijahcilfone/training
0
1import gradio as gr2import os3from pathlib import Path4import argparse5import shutil6from train_dreambooth import run_training7from convertosd import convert8from PIL import Image9from slugify import slugify10import requests11import torch12import zipfile13from diffusers import StableDiffusionPipeline14 15css = '''16 .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}17 .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}18 #component-4, #component-3, #component-10{min-height: 0}19'''20model_to_load = "multimodalart/sd-fine-tunable"21maximum_concepts = 322#Pre download the files even if we don't use it here23StableDiffusionPipeline.from_pretrained(model_to_load)24 25def zipdir(path, ziph):26 # ziph is zipfile handle27 for root, dirs, files in os.walk(path):28 for file in files:29 ziph.write(os.path.join(root, file), 30 os.path.relpath(os.path.join(root, file), 31 os.path.join(path, '..')))32 33def swap_text(option):34 mandatory_liability = "You must have the right to do so and you are liable for the images you use, example:"35 if(option == "object"):36 instance_prompt_example = "cttoy"37 freeze_for = 5038 return [f"You are going to train `object`(s), upload 5-10 images of each object you are planning on training on from different angles/perspectives. {mandatory_liability}:", '''<img src="file/cat-toy.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to 512x512.", freeze_for]39 elif(option == "person"):40 instance_prompt_example = "julcto"41 freeze_for = 10042 return [f"You are going to train a `person`(s), upload 10-20 images of each person you are planning on training on from different angles/perspectives. {mandatory_liability}:", '''<img src="file/person.png" />''', f"You should name the files with a unique word that represent your concept (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to 512x512.", freeze_for]43 elif(option == "style"):44 instance_prompt_example = "trsldamrl"45 freeze_for = 1046 return [f"You are going to train a `style`, upload 10-20 images of the style you are planning on training on. Name the files with the words you would like {mandatory_liability}:", '''<img src="file/trsl_style.png" />''', f"You should name your files with a unique word that represent your concept (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to 512x512.", freeze_for]47 48def count_files(*inputs):49 file_counter = 050 concept_counter = 051 for i, input in enumerate(inputs):52 if(i < maximum_concepts-1):53 files = inputs[i]54 if(files):55 concept_counter+=156 file_counter+=len(files)57 uses_custom = inputs[-1] 58 type_of_thing = inputs[-4]59 if(uses_custom):60 Training_Steps = int(inputs[-3])61 else:62 if(type_of_thing == "person"):63 Training_Steps = file_counter*200*264 else:65 Training_Steps = file_counter*20066 return(gr.update(visible=True, value=f"You are going to train {concept_counter} {type_of_thing}(s), with {file_counter} images for {Training_Steps} steps. This should take around {round(Training_Steps/1.5, 2)} seconds, or {round((Training_Steps/1.5)/3600, 2)} hours. As a reminder, the T4 GPU costs US$0.60 for 1h. Once training is over, don't forget to swap the hardware back to CPU."))67 68def train(*inputs):69 if "IS_SHARED_UI" in os.environ:70 raise gr.Error("This Space only works in duplicated instances")71 if os.path.exists("output_model"): shutil.rmtree('output_model')72 if os.path.exists("instance_images"): shutil.rmtree('instance_images')73 if os.path.exists("diffusers_model.zip"): os.remove("diffusers_model.zip")74 if os.path.exists("model.ckpt"): os.remove("model.ckpt")75 file_counter = 076 for i, input in enumerate(inputs):77 if(i < maximum_concepts-1):78 if(input):79 os.makedirs('instance_images',exist_ok=True)80 files = inputs[i+(maximum_concepts*2)]81 prompt = inputs[i+maximum_concepts]82 if(prompt == "" or prompt == None):83 raise gr.Error("You forgot to define your concept prompt")84 for j, file_temp in enumerate(files):85 file = Image.open(file_temp.name)86 width, height = file.size87 side_length = min(width, height)88 left = (width - side_length)/289 top = (height - side_length)/290 right = (width + side_length)/291 bottom = (height + side_length)/292 image = file.crop((left, top, right, bottom))93 image = image.resize((512, 512))94 extension = file_temp.name.split(".")[1]95 image = image.convert('RGB')96 image.save(f'instance_images/{prompt}_({j+1}).jpg', format="JPEG", quality = 100)97 file_counter += 198 99 os.makedirs('output_model',exist_ok=True)100 uses_custom = inputs[-1] 101 type_of_thing = inputs[-4]102 if(uses_custom):103 Training_Steps = int(inputs[-3])104 Train_text_encoder_for = int(inputs[-2])105 else:106 Training_Steps = file_counter*200107 if(type_of_thing == "object"):108 Train_text_encoder_for=30109 elif(type_of_thing == "person"):110 Train_text_encoder_for=60111 elif(type_of_thing == "style"):112 Train_text_encoder_for=15113 114 class_data_dir = None115 stptxt = int((Training_Steps*Train_text_encoder_for)/100)116 args_general = argparse.Namespace(117 image_captions_filename = True,118 train_text_encoder = True,119 stop_text_encoder_training = stptxt,120 save_n_steps = 0,121 pretrained_model_name_or_path = model_to_load,122 instance_data_dir="instance_images",123 class_data_dir=class_data_dir,124 output_dir="output_model",125 instance_prompt="",126 seed=42,127 resolution=512,128 mixed_precision="fp16",129 train_batch_size=1,130 gradient_accumulation_steps=1,131 use_8bit_adam=True,132 learning_rate=2e-6,133 lr_scheduler="polynomial",134 lr_warmup_steps = 0,135 max_train_steps=Training_Steps, 136 )137 run_training(args_general)138 torch.cuda.empty_cache()139 #convert("output_model", "model.ckpt")140 #shutil.rmtree('instance_images')141 #shutil.make_archive("diffusers_model", 'zip', "output_model")142 with zipfile.ZipFile('diffusers_model.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:143 zipdir('output_model/', zipf)144 torch.cuda.empty_cache()145 return [gr.update(visible=True, value=["diffusers_model.zip"]), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)]146 147def generate(prompt):148 from diffusers import StableDiffusionPipeline149 150 pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)151 pipe = pipe.to("cuda")152 image = pipe(prompt).images[0] 153 return(image)154 155def push(model_name, where_to_upload, hf_token):156 if(not os.path.exists("model.ckpt")):157 convert("output_model", "model.ckpt")158 from huggingface_hub import HfApi, HfFolder, CommitOperationAdd159 from huggingface_hub import create_repo160 model_name_slug = slugify(model_name)161 api = HfApi()162 your_username = api.whoami(token=hf_token)["name"]163 if(where_to_upload == "My personal profile"): 164 model_id = f"{your_username}/{model_name_slug}"165 else:166 model_id = f"sd-dreambooth-library/{model_name_slug}"167 headers = {"Authorization" : f"Bearer: {hf_token}", "Content-Type": "application/json"}168 response = requests.post("https://huggingface.co/organizations/sd-dreambooth-library/share/SSeOwppVCscfTEzFGQaqpfcjukVeNrKNHX", headers=headers)169 170 images_upload = os.listdir("instance_images")171 image_string = ""172 instance_prompt_list = []173 previous_instance_prompt = ''174 for i, image in enumerate(images_upload):175 instance_prompt = image.split("_")[0]176 if(instance_prompt != previous_instance_prompt):177 title_instance_prompt_string = instance_prompt178 instance_prompt_list.append(instance_prompt)179 else:180 title_instance_prompt_string = ''181 previous_instance_prompt = instance_prompt182 image_string = f'''{title_instance_prompt_string}183{image_string}'''184 readme_text = f'''---185license: creativeml-openrail-m186tags:187- text-to-image188---189### {model_name} Dreambooth model trained by {api.whoami(token=hf_token)["name"]} with [Hugging Face Dreambooth Training Space](https://huggingface.co/spaces/multimodalart/dreambooth-training)190 191You run your new concept via `diffusers` [Colab Notebook for Inference](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_dreambooth_inference.ipynb)192 193Sample pictures of this concept:194{image_string}195'''196 #Save the readme to a file197 readme_file = open("README.md", "w")198 readme_file.write(readme_text)199 readme_file.close()200 #Save the token identifier to a file201 text_file = open("token_identifier.txt", "w")202 text_file.write(', '.join(instance_prompt_list))203 text_file.close()204 create_repo(model_id,private=True, token=hf_token)205 operations = [206 CommitOperationAdd(path_in_repo="token_identifier.txt", path_or_fileobj="token_identifier.txt"),207 CommitOperationAdd(path_in_repo="README.md", path_or_fileobj="README.md"),208 CommitOperationAdd(path_in_repo=f"model.ckpt",path_or_fileobj="model.ckpt")209 ]210 api.create_commit(211 repo_id=model_id,212 operations=operations,213 commit_message=f"Upload the model {model_name}",214 token=hf_token215 )216 api.upload_folder(217 folder_path="output_model",218 repo_id=model_id,219 token=hf_token220 )221 api.upload_folder(222 folder_path="instance_images",223 path_in_repo="concept_images",224 repo_id=model_id,225 token=hf_token226 )227 return [gr.update(visible=True, value=f"Successfully uploaded your model. Access it [here](https://huggingface.co/{model_id})"), gr.update(visible=True, value=["diffusers_model.zip", "model.ckpt"])]228 229def convert_to_ckpt():230 convert("output_model", "model.ckpt")231 return gr.update(visible=True, value=["diffusers_model.zip", "model.ckpt"])232 233with gr.Blocks(css=css) as demo:234 with gr.Box():235 if "IS_SHARED_UI" in os.environ:236 gr.HTML('''237 <div class="gr-prose" style="max-width: 80%">238 <h2>Attention - This Space doesn't work in this shared UI</h2>239 <p>For it to work, you have to duplicate the Space and run it on your own profile where a (paid) private GPU will be attributed to it during runtime. As each T4 costs US$0,60/h, it should cost < US$1 to train a model with less than 100 images on default settings!</p> 240 <img class="instruction" src="file/duplicate.png"> 241 <img class="arrow" src="file/arrow.png" />242 </div>243 ''')244 else:245 gr.HTML('''246 <div class="gr-prose" style="max-width: 80%">247 <h2>You have successfully cloned the Dreambooth Training Space</h2>248 <p>If you haven't already, attribute a T4 GPU to it (via the Settings tab) and run the training below. You will be billed by the minute from when you activate the GPU until when you turn it off.</p> 249 </div>250 ''') 251 gr.Markdown("# Dreambooth training")252 gr.Markdown("Customize Stable Diffusion by giving it with few-shot examples. Based on TheLastBen's [fast-DreamBooth Colab](https://colab.research.google.com/github/TheLastBen/fast-stable-diffusion/blob/main/fast-DreamBooth.ipynb) with 🧨 diffusers")253 with gr.Row():254 type_of_thing = gr.Dropdown(label="What would you like to train?", choices=["object", "person", "style"], value="object", interactive=True)255 256 with gr.Row():257 with gr.Column():258 thing_description = gr.Markdown("You are going to train an `object`, upload 5-10 images of the object you are planning on training on from different angles/perspectives. You must have the right to do so and you are liable for the images you use, example:")259 thing_image_example = gr.HTML('''<img src="file/cat-toy.png" />''')260 things_naming = gr.Markdown("You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `cttoy` here). Images will be automatically cropped to 512x512.")261 with gr.Column():262 file_collection = []263 concept_collection = []264 buttons_collection = []265 delete_collection = []266 is_visible = []267 268 row = [None] * maximum_concepts269 for x in range(maximum_concepts):270 ordinal = lambda n: "%d%s" % (n, "tsnrhtdd"[(n // 10 % 10 != 1) * (n % 10 < 4) * n % 10::4])271 if(x == 0):272 visible = True273 is_visible.append(gr.State(value=True))274 else:275 visible = False276 is_visible.append(gr.State(value=False))277 278 file_collection.append(gr.File(label=f"Upload the images for your {ordinal(x+1)} concept", file_count="multiple", interactive=True, visible=visible))279 with gr.Column(visible=visible) as row[x]:280 concept_collection.append(gr.Textbox(label=f"{ordinal(x+1)} concept prompt - use a unique, made up word to avoid collisions")) 281 with gr.Row():282 if(x < maximum_concepts-1):283 buttons_collection.append(gr.Button(value="Add +1 concept", visible=visible))284 if(x > 0):285 delete_collection.append(gr.Button(value=f"Delete {ordinal(x+1)} concept"))286 287 counter_add = 1288 for button in buttons_collection:289 if(counter_add < len(buttons_collection)):290 button.click(lambda:291 [gr.update(visible=True),gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), True, None],292 None, 293 [row[counter_add], file_collection[counter_add], buttons_collection[counter_add-1], buttons_collection[counter_add], is_visible[counter_add], file_collection[counter_add]], queue=False)294 else:295 button.click(lambda:[gr.update(visible=True),gr.update(visible=True), gr.update(visible=False), True], None, [row[counter_add], file_collection[counter_add], buttons_collection[counter_add-1], is_visible[counter_add]], queue=False)296 counter_add += 1297 298 counter_delete = 1299 for delete_button in delete_collection:300 if(counter_delete < len(delete_collection)+1):301 delete_button.click(lambda:[gr.update(visible=False),gr.update(visible=False), gr.update(visible=True), False], None, [file_collection[counter_delete], row[counter_delete], buttons_collection[counter_delete-1], is_visible[counter_delete]], queue=False)302 counter_delete += 1303 304 305 306 with gr.Accordion("Custom Settings", open=False):307 swap_auto_calculated = gr.Checkbox(label="Use custom settings")308 gr.Markdown("If not checked, the number of steps and % of frozen encoder will be tuned automatically according to the amount of images you upload and whether you are training an `object`, `person` or `style` as follows: The number of steps is calculated by number of images uploaded multiplied by 20. The text-encoder is frozen after 10% of the steps for a style, 30% of the steps for an object and is fully trained for persons.")309 steps = gr.Number(label="How many steps", value=800)310 perc_txt_encoder = gr.Number(label="Percentage of the training steps the text-encoder should be trained as well", value=30)311 312 type_of_thing.change(fn=swap_text, inputs=[type_of_thing], outputs=[thing_description, thing_image_example, things_naming, perc_txt_encoder], queue=False)313 training_summary = gr.Textbox("", visible=False, label="Training Summary")314 steps.change(fn=count_files, inputs=file_collection+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[training_summary], queue=False)315 perc_txt_encoder.change(fn=count_files, inputs=file_collection+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[training_summary], queue=False)316 for file in file_collection:317 file.change(fn=count_files, inputs=file_collection+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[training_summary], queue=False)318 train_btn = gr.Button("Start Training")319 with gr.Box(visible=False) as try_your_model:320 gr.Markdown("## Try your model")321 with gr.Row():322 prompt = gr.Textbox(label="Type your prompt")323 result_image = gr.Image()324 generate_button = gr.Button("Generate Image")325 with gr.Box(visible=False) as push_to_hub:326 gr.Markdown("## Push to Hugging Face Hub")327 model_name = gr.Textbox(label="Name of your model", placeholder="Tarsila do Amaral Style")328 where_to_upload = gr.Dropdown(["My personal profile", "Public Library"], label="Upload to")329 gr.Markdown("[A Hugging Face write access token](https://huggingface.co/settings/tokens), go to \"New token\" -> Role : Write. A regular read token won't work here.")330 hf_token = gr.Textbox(label="Hugging Face Write Token")331 push_button = gr.Button("Push to the Hub")332 result = gr.File(label="Download the uploaded models in the diffusers format", visible=True)333 success_message_upload = gr.Markdown(visible=False)334 convert_button = gr.Button("Convert to CKPT", visible=False)335 336 train_btn.click(fn=train, inputs=is_visible+concept_collection+file_collection+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[result, try_your_model, push_to_hub, convert_button])337 generate_button.click(fn=generate, inputs=prompt, outputs=result_image)338 push_button.click(fn=push, inputs=[model_name, where_to_upload, hf_token], outputs=[success_message_upload, result])339 convert_button.click(fn=convert_to_ckpt, inputs=[], outputs=result)340demo.launch(debug=True)