127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
from argparse import ArgumentParser
|
||
|
|
import time
|
||
|
|
from time import gmtime, strftime
|
||
|
|
from PIL import Image
|
||
|
|
import torch
|
||
|
|
from torch import autocast
|
||
|
|
from tqdm.auto import tqdm
|
||
|
|
import numpy as np
|
||
|
|
import piexif
|
||
|
|
import piexif.helper
|
||
|
|
|
||
|
|
from diffusers import (
|
||
|
|
DDIMScheduler,
|
||
|
|
#PNDMScheduler,
|
||
|
|
)
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
##Generating with stable diffusion
|
||
|
|
device = "cuda"
|
||
|
|
model_path = "CompVis/stable-diffusion-v1-4"
|
||
|
|
|
||
|
|
# Using DDIMScheduler as anexample,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()
|
||
|
|
|
||
|
|
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 = []
|
||
|
|
|
||
|
|
startStrength = 0.86
|
||
|
|
deltaStrength = 0.02
|
||
|
|
endStrength = 0.941
|
||
|
|
|
||
|
|
startScale = 7.5
|
||
|
|
deltaScale = 2.5
|
||
|
|
endScale = 15.0
|
||
|
|
|
||
|
|
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))
|
||
|
|
|
||
|
|
print("Found %i pictures", len(path_img))
|
||
|
|
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))
|
||
|
|
|
||
|
|
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)
|
||
|
|
strenght = startStrength
|
||
|
|
while strenght <= endStrength:
|
||
|
|
guidance_scale = startScale
|
||
|
|
while guidance_scale <= endScale:
|
||
|
|
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")
|
||
|
|
|
||
|
|
allwork+=1
|
||
|
|
|
||
|
|
guidance_scale+=deltaScale
|
||
|
|
strenght+=deltaStrength
|
||
|
|
|
||
|
|
counterr+=1
|
||
|
|
|
||
|
|
print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds")
|