135 lines
3.9 KiB
Python
135 lines
3.9 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import time
|
|
from time import strftime
|
|
from PIL import Image
|
|
import torch
|
|
from torch import autocast
|
|
import numpy as np
|
|
import piexif
|
|
import piexif.helper
|
|
|
|
from diffusers import (
|
|
#DDIMScheduler,
|
|
PNDMScheduler,
|
|
StableDiffusionImg2ImgPipeline
|
|
)
|
|
|
|
t = torch.cuda.get_device_properties(0).total_memory
|
|
if t <= 6400000000:
|
|
print("Running with less than 6gb memory. Working is not garanted")
|
|
|
|
file1 = open("prompt.txt", "r+")
|
|
text = file1.read()
|
|
print(text)
|
|
|
|
# Generating with stable diffusion
|
|
device = "cuda"
|
|
model_path = "CompVis/stable-diffusion-v1-4"
|
|
|
|
scheduler = PNDMScheduler.from_config(model_path, subfolder="scheduler")
|
|
#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
|
|
).to(device)
|
|
pipe.enable_attention_slicing()
|
|
|
|
|
|
def dummy_checker(images, **kwargs):
|
|
return images, False
|
|
|
|
|
|
pipe.safety_checker = dummy_checker
|
|
|
|
|
|
def preprocess(image):
|
|
w, h = image.size
|
|
# resize to integer multiple of 32
|
|
w, h = map(lambda x: x - x % 32, (w, h))
|
|
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 " + str(len(path_img)) + " pictures")
|
|
start_time = time.time()
|
|
|
|
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
|
|
|
|
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).images[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
|
|
|
|
strenght += deltaStrength
|
|
guidance_scale += deltaScale
|
|
|
|
counterr += 1
|
|
|
|
print("Made " + str(allwork) + " pictures in " +
|
|
str(time.time()-start_time) + " seconds")
|