diff --git a/README.md b/README.md index 8dfb7b2..d4eed73 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,9 @@ c. Create a write token/Copy an existing Token key and enter in the cmd window. 5. You can find output files in .\\data\\out\\%Timestamp% folder. Jpeg contain commentaries in metadata with some settings used while generating. ## Requirements -* 11GB of free space. -* Nvidia card with 8GB+ video memory. - * Currently min dalle can run in 6gb. Testing purpose -## Additional info -Tested on RTX3070. One picture was making 12 - 14 seconds. \ No newline at end of file +* 10GB of free space. +* Nvidia card with 6GB+ video memory. + +## Tests +RTX3070: 12 - 14 seconds per picture. +RTX2060_Laptop: 21 - 25 seconds per picture. \ No newline at end of file diff --git a/scripts/stable_diff.py b/scripts/stable_diff.py index 2b6c9cf..08a1cbe 100644 --- a/scripts/stable_diff.py +++ b/scripts/stable_diff.py @@ -20,113 +20,114 @@ from diffusers import ( from diffusers import StableDiffusionImg2ImgPipeline t = torch.cuda.get_device_properties(0).total_memory -if t <=8500000000: - print("Not enough GPU memory to generate pictures") -else: - file1 = open("prompt.txt","r+") - text = file1.read() - print(text) +if t<=6400000000: + print("Running with less than 6gb memory. Working is not garanted") - ##Generating with stable diffusion - device = "cuda" - model_path = "CompVis/stable-diffusion-v1-4" +file1 = open("prompt.txt","r+") +text = file1.read() +print(text) - # Using DDIMScheduler as anexample,this also works with PNDMScheduler +##Generating with stable diffusion +device = "cuda" +model_path = "CompVis/stable-diffusion-v1-4" + + # Using DDIMScheduler as an example,this also works with PNDMScheduler # uncomment this line if you want to use it. #scheduler = PNDMScheduler.from_config(model_path, subfolder="scheduler", use_auth_token=True) - scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False) - pipe = StableDiffusionImg2ImgPipeline.from_pretrained( - model_path, - scheduler=scheduler, - revision="fp16", - torch_dtype=torch.float16, - use_auth_token=True - ).to(device) - if t <= 10500000000: - pipe.enable_attention_slicing() +scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False) +pipe = StableDiffusionImg2ImgPipeline.from_pretrained( + model_path, + scheduler=scheduler, + revision="fp16", + torch_dtype=torch.float16, + use_auth_token=True +).to(device) +if t <= 11000000000: + print("Less then 11gb video memory. Running with attention slicing.") + pipe.enable_attention_slicing() - def dummy_checker(images, **kwargs): - return images, False - pipe.safety_checker = dummy_checker +def dummy_checker(images, **kwargs): + return images, False +pipe.safety_checker = dummy_checker - def preprocess(image): - w, h = image.size - w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), Image.Resampling.LANCZOS) - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - image = torch.from_numpy(image) - return 2.*image - 1. +def preprocess(image): + w, h = image.size + w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 + image = image.resize((w, h), Image.Resampling.LANCZOS) + image = np.array(image).astype(np.float32) / 255.0 + image = image[None].transpose(0, 3, 1, 2) + image = torch.from_numpy(image) + return 2.*image - 1. - path_img = [] +path_img = [] - startStrength = 0.86 - deltaStrength = 0.02 - endStrength = 0.941 +startStrength = 0.86 +deltaStrength = 0.02 +endStrength = 0.941 - startScale = 7.5 - deltaScale = 2.5 - endScale = 15.0 +startScale = 7.5 +deltaScale = 2.5 +endScale = 15.0 - startSeed = 1022 - endSeed = 1024 +startSeed = 1022 +endSeed = 1024 - directory_in = "./data/input" - if not os.path.exists(directory_in): - os.makedirs(directory_in) - for root, subdirectories, files in os.walk(directory_in): - for filename in files: - if filename.endswith(".png"): - path_img.append(os.path.join(root, filename)) +directory_in = "./data/input" +if not os.path.exists(directory_in): + os.makedirs(directory_in) +for root, subdirectories, files in os.walk(directory_in): + for filename in files: + if filename.endswith(".png"): + path_img.append(os.path.join(root, filename)) - print("Found " +str(len(path_img)) +" pictures") - start_time = time.time() +print("Found " +str(len(path_img)) +" pictures") +start_time = time.time() - counterr=0 - allwork=0 - directory="./data/out/" + strftime("%Y-%m-%d_%H-%M-%S", gmtime()) + "/" - if not os.path.exists(directory): - os.makedirs(directory) - with open(directory+"prompt.txt", 'w') as f: - f.write(text) - shutil.copytree(directory_in, directory+"input") - for i in path_img: - print(i) - if not os.path.exists(directory+str(counterr)): - os.makedirs(directory+str(counterr)) +counterr=0 +allwork=0 +directory="./data/out/" + strftime("%Y-%m-%d_%H-%M-%S") + "/" +if not os.path.exists(directory): + os.makedirs(directory) +with open(directory+"prompt.txt", 'w') as f: + f.write(text) +shutil.copytree(directory_in, directory+"input") +for i in path_img: + print(i) + if not os.path.exists(directory+str(counterr)): + os.makedirs(directory+str(counterr)) - init_img = Image.open(i) - init_img = init_img.resize((768, 512)) - init_image = preprocess(init_img) - prompt = text + init_img = Image.open(i) + init_img = init_img.resize((768, 512)) + init_image = preprocess(init_img) + prompt = text - for seed in range(endSeed-startSeed+1): - generator = torch.Generator(device=device).manual_seed(startSeed+seed) - guidance_scale = startScale - while guidance_scale <= endScale: - strenght = startStrength - while strenght <= endStrength: - with autocast("cuda"): - image = pipe(prompt=prompt, init_image=init_image, strength=strenght, guidance_scale=guidance_scale, generator=generator)["sample"][0] - image.save(directory+str(counterr)+"/" + str(allwork) +".jpg") + for seed in range(endSeed-startSeed+1): + generator = torch.Generator(device=device).manual_seed(startSeed+seed) + guidance_scale = startScale + while guidance_scale <= endScale: + strenght = startStrength + while strenght <= endStrength: + with autocast("cuda"): + image = pipe(prompt=prompt, init_image=init_image, strength=strenght, guidance_scale=guidance_scale, generator=generator)["sample"][0] + image.save(directory+str(counterr)+"/" + str(allwork) +".jpg") - exif_dict = piexif.load(directory+str(counterr)+"/" + str(allwork) +".jpg") - userCommentAsDict = {} - userCommentAsDict['Seed']=str(seed+startSeed) - userCommentAsDict['Strength']=str(strenght) - userCommentAsDict['Guidance_scale']=str(guidance_scale) - user_comment = piexif.helper.UserComment.dump(json.dumps(userCommentAsDict)) - exif_dict["Exif"][piexif.ExifIFD.UserComment] = user_comment - exif_bytes = piexif.dump(exif_dict) - piexif.insert(exif_bytes, directory+str(counterr)+"/" + str(allwork) +".jpg") + exif_dict = piexif.load(directory+str(counterr)+"/" + str(allwork) +".jpg") + userCommentAsDict = {} + userCommentAsDict['Seed']=str(seed+startSeed) + userCommentAsDict['Strength']=str(strenght) + userCommentAsDict['Guidance_scale']=str(guidance_scale) + user_comment = piexif.helper.UserComment.dump(json.dumps(userCommentAsDict)) + exif_dict["Exif"][piexif.ExifIFD.UserComment] = user_comment + exif_bytes = piexif.dump(exif_dict) + piexif.insert(exif_bytes, directory+str(counterr)+"/" + str(allwork) +".jpg") - allwork+=1 + allwork+=1 - strenght+=deltaStrength - guidance_scale+=deltaScale + strenght+=deltaStrength + guidance_scale+=deltaScale - counterr+=1 + counterr+=1 - print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds") \ No newline at end of file +print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds") \ No newline at end of file