app.py

import torch
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance
from datetime import datetime, timedelta
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoencoderTiny
from diffusers.models.attention_processor import SlicedAttnProcessor
from tqdm import tqdm
from pathlib import Path

def adjust_gamma(img, gamma=0.4):
    npim = np.array(img)
    npim_gamma = 255.0 * (npim / 255.0) ** gamma
    return Image.fromarray(np.uint8(npim_gamma))

atkbold = ImageFont.truetype("Atkinson-Hyperlegible-Bold-102.otf", 380)
atkbold_smol = ImageFont.truetype("Atkinson-Hyperlegible-Bold-102.otf", 40)

image_size = (1600, 900)
screen_size = image_size

def mask_image(timestamp):
    mask_text = timestamp.strftime(f"%H\u2009%M\u2009%S")
    time_img = Image.new("L", image_size, (0,))
    draw = ImageDraw.Draw(time_img)
    draw.multiline_text(
        xy=(-30,120),
        text=mask_text,
        fill=(255,),
        font=atkbold,
        align="center",
        spacing=-10,
    )
    # return time_img
    (i_left, i_top, i_right, i_bottom) = time_img.getbbox()
    # pad the image horizonally to the full size
    i_left = 0
    i_right = image_size[0]
    time_img = time_img.crop((i_left, i_top, i_right, i_bottom))
    return ImageOps.pad(time_img, image_size)

preferred_dtype = torch.float32
preferred_device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
#preferred_device = "cpu"

controlnet = ControlNetModel.from_pretrained(
    "monster-labs/control_v1p_sd15_qrcode_monster",
    subfolder="v2",
    torch_dtype=preferred_dtype,
).to(preferred_device)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "SimianLuo/LCM_Dreamshaper_v7",
    controlnet=controlnet,
    torch_dtype=preferred_dtype,
    safety_checker=None,
).to(preferred_device)

#pipe.vae.set_attn_processor(SlicedAttnProcessor(4))
#pipe.unet.set_attn_processor(SlicedAttnProcessor(4))

current_denoising_steps = 4
current_latency = 0
rounding_minutes = 1
target_filename = "beauty.png"
mask_image(timestamp=datetime.now()).save(target_filename)

cali1 = "desert landscape with tall mountains and cactus and boulders at sunrise with the sun on the horizon"
cali2 = "stony river in a sunny redwood forest with salmon and deer and bears and mushrooms"
cali3 = "beach, tall cliffs, sun at the horizon, albatross eating fish, no one on the beach, boulders and tide pools in the shallow water"
cali4 = "nighttime photo of a desert landscape with the milky way in the sky and boulders on a shallow lake bed surrounded by tall mountains"

prompts = [
#    cali1,
#    cali2,
#    cali3,
    cali4,
]
conditioning_scales = {
    cali1: 0.45,
    cali2: 0.5,
    cali3: 0.45,
    cali4: 0.6,
}
negative_prompt = "low quality, ugly, wrong"

xdalidemo = True
fps = 10
if xdalidemo:
    iteration_range = range(86400 * fps)
else:
    iteration_range = range(86400 * 365 * 80)

def ease(x):
    return (x * x) * (3 - 2 * x) # (x * x * x * x * x) * (126 - 420 * x + 540 * x * x - 315 * x * x * x + 70 * x * x * x * x)

for iteration in tqdm(iteration_range):
    if xdalidemo:
        synthetic_time = datetime(year=2000,month=1,day=1,hour=0,minute=0,second=0) + timedelta(microseconds=iteration*(1000000/fps))
        target_filename = f"face-render-{iteration:08}.png"
        current_latency = 0
        pre_render_time = synthetic_time
        target_time_plus_midpoint = synthetic_time
        rounded_target_time = synthetic_time
        this_second = synthetic_time
        next_second = synthetic_time + timedelta(seconds=1)
        this_mask = mask_image(this_second)
        next_mask = mask_image(next_second)
        easing_step = ease((iteration % fps) / fps)
        current_mask_image = Image.fromarray(np.array(this_mask) * (1-easing_step) + np.array(next_mask) * (easing_step))
    else:
        pre_render_time = datetime.now()
        target_time_plus_midpoint = pre_render_time + timedelta(seconds=(current_latency + rounding_minutes * 60 / 2))
        rounded_target_time = target_time_plus_midpoint - timedelta(minutes=target_time_plus_midpoint.minute - target_time_plus_midpoint.minute // rounding_minutes * rounding_minutes)
        current_mask_image = mask_image(timestamp=rounded_target_time)

    print(f"current_latency: {current_latency}, pre_render_time: {pre_render_time}, rounded_target_time: {rounded_target_time}, current_denoising_steps: {current_denoising_steps}\n")

    if Path(target_filename).exists():
        continue


#    image = current_mask_image.convert("RGB")
    image = pipe(
        prompt=prompts[iteration % len(prompts)],
        negative_prompt=negative_prompt,
        image=current_mask_image,
        num_inference_steps=current_denoising_steps,
        guidance_scale=7.0,
        controlnet_conditioning_scale=conditioning_scales[prompts[iteration % len(prompts)]],
        generator=torch.manual_seed(int(rounded_target_time.timestamp()) // 60),
        height=image_size[1],
        width=image_size[0],
    ).images[0]
    #image = adjust_gamma(image, gamma=0.5)
    #image = ImageEnhance.Sharpness(image).enhance(5)
    #image = image.resize(screen_size)
    if True:
        draw = ImageDraw.Draw(image)
        draw.text((60, screen_size[1]-60), f"leebutterman.com", fill=(255,255,255), font=atkbold_smol)
    image.save(target_filename)
    post_render_time = datetime.now()
    current_latency = post_render_time.timestamp() - pre_render_time.timestamp()