You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
import os.path
|
|
|
|
import torch
|
|
from PIL import ImageDraw
|
|
|
|
from imaginairy import ImaginePrompt, LazyLoadingImage, imagine, imagine_image_files
|
|
from imaginairy.api import load_model
|
|
from imaginairy.img_log import ImageLoggingContext, filesafe_text, log_latent
|
|
from imaginairy.modules.clip_embedders import FrozenCLIPEmbedder
|
|
from imaginairy.samplers.ddim import DDIMSampler
|
|
from imaginairy.utils import get_device, pillow_img_to_torch_image
|
|
from tests import TESTS_FOLDER
|
|
|
|
|
|
def experiment_text_conditioning_combos():
|
|
"""
|
|
Can we do math with the embeddings?
|
|
|
|
Yes. it works but doesn't look great.
|
|
"""
|
|
embedder = FrozenCLIPEmbedder()
|
|
embedder.to(get_device())
|
|
|
|
beach_e = embedder.encode(["a beach"])
|
|
beach_water_e = embedder.encode(["a beach. ocean, waves, water"])
|
|
waterness = beach_water_e - beach_e
|
|
waterless_beach = beach_e - waterness
|
|
|
|
imagine_image_files(
|
|
[ImaginePrompt("waterless_beach", conditioning=waterless_beach, seed=1)],
|
|
outdir=f"{TESTS_FOLDER}/test_output",
|
|
)
|
|
imagine_image_files(
|
|
[ImaginePrompt("waterness", conditioning=waterness, seed=1)],
|
|
outdir=f"{TESTS_FOLDER}/test_output",
|
|
)
|
|
imagine_image_files(
|
|
[ImaginePrompt("beach", conditioning=beach_e, seed=1)],
|
|
outdir=f"{TESTS_FOLDER}/test_output",
|
|
)
|
|
|
|
|
|
def experiment_step_repeats():
|
|
"""
|
|
Run the same step over and over on an image without noise
|
|
|
|
Removes detail from the image.
|
|
"""
|
|
model = load_model()
|
|
model.to(get_device())
|
|
model.eval()
|
|
embedder = FrozenCLIPEmbedder()
|
|
embedder.to(get_device())
|
|
|
|
sampler = DDIMSampler(model)
|
|
sampler.make_schedule(1000)
|
|
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg")
|
|
init_image, _, h = pillow_img_to_torch_image(
|
|
img,
|
|
max_height=512,
|
|
max_width=512,
|
|
)
|
|
init_image = init_image.to(get_device())
|
|
init_latent = model.get_first_stage_encoding(model.encode_first_stage(init_image))
|
|
log_latent(init_latent, "init_latent")
|
|
|
|
base_count = 1
|
|
neutral_embedding = embedder.encode([""])
|
|
outdir = f"{TESTS_FOLDER}/test_output"
|
|
|
|
def _record_step(img, description, step_count, prompt):
|
|
steps_path = os.path.join(outdir, "steps", f"{base_count:08}_S{prompt.seed}")
|
|
os.makedirs(steps_path, exist_ok=True)
|
|
filename = f"{base_count:08}_S{prompt.seed}_step{step_count:04}_{filesafe_text(description)[:40]}.png"
|
|
destination = os.path.join(steps_path, filename)
|
|
draw = ImageDraw.Draw(img)
|
|
draw.text((10, 10), str(description))
|
|
img.save(destination)
|
|
|
|
with ImageLoggingContext(
|
|
prompt=ImaginePrompt(""),
|
|
model=model,
|
|
img_callback=_record_step,
|
|
):
|
|
x_prev = init_latent
|
|
index = 50
|
|
base_count = index
|
|
t = torch.Tensor([index]).to(get_device())
|
|
# noise_pred = model.apply_model(init_latent, t, neutral_embedding)
|
|
# log_latent(noise_pred, "noise prediction")
|
|
for _ in range(100):
|
|
x_prev, pred_x0 = sampler.p_sample_ddim(x_prev, neutral_embedding, t, index)
|
|
log_latent(pred_x0, "pred_x0")
|
|
x_prev = pred_x0
|
|
|
|
|
|
def experiment_repeated_img_2_img():
|
|
"""
|
|
Experiment with putting an image repeatedly through image2image
|
|
|
|
It creates screwy images
|
|
"""
|
|
outdir = f"{TESTS_FOLDER}/test_output/img2img2img"
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg")
|
|
img.save(f"{outdir}/0.png")
|
|
for step_num in range(50):
|
|
prompt = ImaginePrompt(
|
|
"Beach at Sainte Adresse. hyperealistic photo. sharp focus, canon 5d",
|
|
init_image=img,
|
|
init_image_strength=0.50,
|
|
width=512,
|
|
height=512,
|
|
steps=50,
|
|
sampler_type="DDIM",
|
|
)
|
|
|
|
result = next(imagine(prompt))
|
|
img = result.img
|
|
os.makedirs(outdir, exist_ok=True)
|
|
img.save(f"{outdir}/{step_num:04}.png")
|