Sudharma/GenerativeAI
0
1# -*- coding: utf-8 -*-2"""AI_Generative.ipynb3 4Automatically generated by Colaboratory.5 6Original file is located at7 https://colab.research.google.com/drive/1TlivRo6A5WTBKbAWmNoYPblVzjNXLoo68 9## Install Requirements10"""11 12# Commented out IPython magic to ensure Python compatibility.13!wget -q https://github.com/ShivamShrirao/diffusers/raw/main/examples/dreambooth/train_dreambooth.py14!wget -q https://github.com/ShivamShrirao/diffusers/raw/main/scripts/convert_diffusers_to_original_stable_diffusion.py15# %pip install -qq git+https://github.com/ShivamShrirao/diffusers16# %pip install -q -U --pre triton17# %pip install -q accelerate transformers ftfy bitsandbytes==0.35.0 gradio natsort safetensors xformers18 19!apt-get install wget20!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip21!unzip ngrok-stable-linux-amd64.zip22 23!mkdir -p ~/.huggingface24HUGGINGFACE_TOKEN = "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"25!echo -n "{HUGGINGFACE_TOKEN}" > ~/.huggingface/token26 27"""## Settings and run"""28 29save_to_gdrive = True 30if save_to_gdrive:31 from google.colab import drive32 drive.mount('/content/drive')33MODEL_NAME = "runwayml/stable-diffusion-v1-5"34OUTPUT_DIR = "stable_diffusion_weights/cxyz" 35if save_to_gdrive:36 OUTPUT_DIR = "/content/drive/MyDrive/" + OUTPUT_DIR37else:38 OUTPUT_DIR = "/content/" + OUTPUT_DIR39 40print(f"[*] Weights will be saved at {OUTPUT_DIR}")41 42!mkdir -p $OUTPUT_DIR43 44"""# Start Training45 46Use the table below to choose the best flags based on your memory and speed requirements. Tested on Tesla T4 GPU.47 48 49| `fp16` | `train_batch_size` | `gradient_accumulation_steps` | `gradient_checkpointing` | `use_8bit_adam` | GB VRAM usage | Speed (it/s) |50| ---- | ------------------ | ----------------------------- | ----------------------- | --------------- | ---------- | ------------ |51| fp16 | 1 | 1 | TRUE | TRUE | 9.92 | 0.93 |52| no | 1 | 1 | TRUE | TRUE | 10.08 | 0.42 |53| fp16 | 2 | 1 | TRUE | TRUE | 10.4 | 0.66 |54| fp16 | 1 | 1 | FALSE | TRUE | 11.17 | 1.14 |55| no | 1 | 1 | FALSE | TRUE | 11.17 | 0.49 |56| fp16 | 1 | 2 | TRUE | TRUE | 11.56 | 1 |57| fp16 | 2 | 1 | FALSE | TRUE | 13.67 | 0.82 |58| fp16 | 1 | 2 | FALSE | TRUE | 13.7 | 0.83 |59| fp16 | 1 | 1 | TRUE | FALSE | 15.79 | 0.77 |60 61Add `--gradient_checkpointing` flag for around 9.92 GB VRAM usage.62 63remove `--use_8bit_adam` flag for full precision. Requires 15.79 GB with `--gradient_checkpointing` else 17.8 GB.64 65remove `--train_text_encoder` flag to reduce memory usage further, degrades output quality.66"""67 68concepts_list = [69 {70 "instance_prompt": "cxyz",71 "class_prompt": "human image",72 "instance_data_dir": "/content/drive/MyDrive/myimage",73 "class_data_dir": "/content/drive/MyDrive/human pics"74 }75 ]76import json77import os78for c in concepts_list:79 os.makedirs(c["instance_data_dir"], exist_ok=True)80 81with open("concepts_list.json", "w") as f:82 json.dump(concepts_list, f, indent=4)83 84import os85from google.colab import files86import shutil87 88for c in concepts_list:89 print(f"Uploading instance images for `{c['instance_prompt']}`")90 uploaded = files.upload()91 for filename in uploaded.keys():92 dst_path = os.path.join(c['instance_data_dir'], filename)93 shutil.move(filename, dst_path)94 95!accelerate launch train_dreambooth.py \96 --pretrained_model_name_or_path=$MODEL_NAME \97 --pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse" \98 --output_dir=$OUTPUT_DIR \99 --revision="fp16" \100 --with_prior_preservation --prior_loss_weight=1.0 \101 --seed=1337 \102 --resolution=512 \103 --train_batch_size=1 \104 --train_text_encoder \105 --mixed_precision="fp16" \106 --use_8bit_adam \107 --gradient_accumulation_steps=1 \108 --learning_rate=1e-6 \109 --lr_scheduler="constant" \110 --lr_warmup_steps=168 \111 --num_class_images=252 \112 --sample_batch_size=4 \113 --max_train_steps=1680 \114 --save_interval=10000 \115 --save_sample_prompt="cxyz" \116 --concepts_list="concepts_list.json"117 118WEIGHTS_DIR = "" 119if WEIGHTS_DIR == "":120 from natsort import natsorted121 from glob import glob122 import os123 WEIGHTS_DIR = natsorted(glob(OUTPUT_DIR + os.sep + "*"))[-1]124print(f"[*] WEIGHTS_DIR={WEIGHTS_DIR}")125 126import os127import matplotlib.pyplot as plt128import matplotlib.image as mpimg129 130weights_folder = OUTPUT_DIR131folders = sorted([f for f in os.listdir(weights_folder) if f != "0"], key=lambda x: int(x))132 133row = len(folders)134col = len(os.listdir(os.path.join(weights_folder, folders[0], "samples")))135scale = 4136fig, axes = plt.subplots(row, col, figsize=(col*scale, row*scale), gridspec_kw={'hspace': 0, 'wspace': 0})137 138for i, folder in enumerate(folders):139 folder_path = os.path.join(weights_folder, folder)140 image_folder = os.path.join(folder_path, "samples")141 images = [f for f in os.listdir(image_folder)]142 for j, image in enumerate(images):143 if row == 1:144 currAxes = axes[j]145 else:146 currAxes = axes[i, j]147 if i == 0:148 currAxes.set_title(f"Image {j}")149 if j == 0:150 currAxes.text(-0.1, 0.5, folder, rotation=0, va='center', ha='center', transform=currAxes.transAxes)151 image_path = os.path.join(image_folder, image)152 img = mpimg.imread(image_path)153 currAxes.imshow(img, cmap='gray')154 currAxes.axis('off')155 156plt.tight_layout()157plt.savefig('grid.png', dpi=72)158 159!./ngrok authtoken 2N8CgXpCOPUYLXNVNIy6KNCV5nD_3Z735wg5PANop9cKoxyYN160 161"""## Convert weights to ckpt to use in web UIs like AUTOMATIC1111.162 163## Inference164"""165 166import torch167!pip install flask_ngrok168from torch import autocast169from diffusers import StableDiffusionPipeline, DDIMScheduler170from IPython.display import display171from flask import Flask172 173app = Flask(__name__)174model_path = WEIGHTS_DIR 175pipe = StableDiffusionPipeline.from_pretrained(model_path, safety_checker=None, torch_dtype=torch.float16).to("cuda")176pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)177pipe.enable_xformers_memory_efficient_attention()178g_cuda = None179 180g_cuda = torch.Generator(device='cuda')181seed = 52362 182g_cuda.manual_seed(seed)183 184from flask import Flask185from flask import request186from flask import Response187from flask_ngrok import run_with_ngrok188import os189import json190 191 192 193app = Flask(__name__)194run_with_ngrok(app)195 196if not os.path.exists('images'):197 os.makedirs('images')198 199@app.route("/test")200def test():201 return "Vanakam"202 203@app.route("/avatar-face")204def generateImage():205 prompt = request.form['prompt']206 negative_prompt = request.form['negative_prompt']207 num_samples = request.form['num_samples']208 guidance_scale = 7.9209 num_inference_steps = 50 210 height = 512211 width = 512 212 images_list = []213 with autocast("cuda"), torch.inference_mode():214 images = pipe(215 prompt,216 height=height,217 width=width,218 negative_prompt=negative_prompt,219 num_images_per_prompt=num_samples,220 num_inference_steps=num_inference_steps,221 guidance_scale=guidance_scale,222 generator=g_cuda223 ).images224 225 for img in images:226 os.rename(os.getcwd()+"/images",img)227 images_list.append("http://eb12-34-141-227-3.ngrok.io/"+os.getcwd()+"/images"+img)228 #display(img)229 return Response(json.dumps(images_list), mimetype='application/json')230app.run()231 232#@title Free runtime memory233exit()