Bryce
2 years ago
parent
47c6bcee59
commit
66c640ce7b
@ -0,0 +1,70 @@
|
||||
import hashlib
|
||||
import random
|
||||
|
||||
import numpy
|
||||
|
||||
|
||||
class WeightedPrompt:
|
||||
def __init__(self, text, weight=1):
|
||||
self.text = text
|
||||
self.weight = weight
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.weight}*({self.text})"
|
||||
|
||||
|
||||
class ImaginePrompt:
|
||||
def __init__(
|
||||
self,
|
||||
prompt=None,
|
||||
seed=None,
|
||||
prompt_strength=7.5,
|
||||
sampler_type="PLMS",
|
||||
init_image=None,
|
||||
init_image_strength=0.3,
|
||||
steps=50,
|
||||
height=512,
|
||||
width=512,
|
||||
upscale=False,
|
||||
fix_faces=False,
|
||||
parts=None,
|
||||
):
|
||||
prompt = prompt if prompt is not None else "a scenic landscape"
|
||||
if isinstance(prompt, str):
|
||||
self.prompts = [WeightedPrompt(prompt, 1)]
|
||||
else:
|
||||
self.prompts = prompt
|
||||
self.init_image = init_image
|
||||
self.init_image_strength = init_image_strength
|
||||
self.prompts.sort(key=lambda p: p.weight, reverse=True)
|
||||
self.seed = random.randint(1, 1_000_000_000) if seed is None else seed
|
||||
self.prompt_strength = prompt_strength
|
||||
self.sampler_type = sampler_type
|
||||
self.steps = steps
|
||||
self.height = height
|
||||
self.width = width
|
||||
self.upscale = upscale
|
||||
self.fix_faces = fix_faces
|
||||
self.parts = parts or {}
|
||||
|
||||
@property
|
||||
def prompt_text(self):
|
||||
if len(self.prompts) == 1:
|
||||
return self.prompts[0].text
|
||||
return "|".join(str(p) for p in self.prompts)
|
||||
|
||||
|
||||
class ImagineResult:
|
||||
def __init__(self, img, prompt):
|
||||
self.img = img
|
||||
self.prompt = prompt
|
||||
|
||||
def cv2_img(self):
|
||||
open_cv_image = numpy.array(self.img)
|
||||
# Convert RGB to BGR
|
||||
open_cv_image = open_cv_image[:, :, ::-1].copy()
|
||||
return open_cv_image
|
||||
# return cv2.cvtColor(numpy.array(self.img), cv2.COLOR_RGB2BGR)
|
||||
|
||||
def md5(self):
|
||||
return hashlib.md5(self.img.tobytes()).hexdigest()
|
||||
@ -0,0 +1 @@
|
||||
pytest
|
||||
@ -0,0 +1,3 @@
|
||||
import os.path
|
||||
|
||||
TESTS_FOLDER = os.path.dirname(__file__)
|
||||
@ -0,0 +1,60 @@
|
||||
from imaginairy.imagine import imagine_images, imagine_image_files
|
||||
from imaginairy.schema import ImaginePrompt, WeightedPrompt
|
||||
from . import TESTS_FOLDER
|
||||
|
||||
|
||||
def test_imagine():
|
||||
prompt = ImaginePrompt("a scenic landscape", width=512, height=256, steps=20, seed=1)
|
||||
result = next(imagine_images(prompt))
|
||||
assert result.md5() == '4c5957c498881d365cfcf13014812af0'
|
||||
result.img.save(f"{TESTS_FOLDER}/test_output/scenic_landscape.png")
|
||||
|
||||
|
||||
def test_img_to_img():
|
||||
prompt = ImaginePrompt(
|
||||
"a photo of a beach",
|
||||
init_image=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg",
|
||||
init_image_strength=0.5,
|
||||
width=512,
|
||||
height=512,
|
||||
steps=50,
|
||||
seed=1,
|
||||
sampler_type="DDIM",
|
||||
)
|
||||
out_folder = f"{TESTS_FOLDER}/test_output"
|
||||
out_folder = '/home/bryce/Mounts/drennanfiles/art/tests'
|
||||
imagine_image_files(prompt, outdir=out_folder)
|
||||
|
||||
|
||||
def test_img_to_file():
|
||||
prompt = ImaginePrompt(
|
||||
[WeightedPrompt("an old growth forest, diffuse light poking through the canopy. high-resolution, nature photography, nat geo photo")],
|
||||
# init_image=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg",
|
||||
init_image_strength=0.5,
|
||||
width=512+64,
|
||||
height=512-64,
|
||||
steps=50,
|
||||
# seed=2,
|
||||
sampler_type="PLMS",
|
||||
upscale=True
|
||||
)
|
||||
out_folder = f"{TESTS_FOLDER}/test_output"
|
||||
out_folder = '/home/bryce/Mounts/drennanfiles/art/tests'
|
||||
imagine_image_files(prompt, outdir=out_folder)
|
||||
|
||||
|
||||
def test_img_conditioning():
|
||||
prompt = ImaginePrompt(
|
||||
"photo",
|
||||
init_image=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg",
|
||||
init_image_strength=0.5,
|
||||
width=512+64,
|
||||
height=512-64,
|
||||
steps=50,
|
||||
# seed=2,
|
||||
sampler_type="PLMS",
|
||||
upscale=True
|
||||
)
|
||||
out_folder = f"{TESTS_FOLDER}/test_output"
|
||||
out_folder = '/home/bryce/Mounts/drennanfiles/art/tests'
|
||||
imagine_image_files(prompt, outdir=out_folder, record_steps=True)
|
||||
Loading…
Reference in New Issue