Spaces:
Sleeping
Sleeping
DEVICE = 'cpu' | |
import gradio as gr | |
import numpy as np | |
from sklearn.svm import LinearSVC | |
from sklearn import preprocessing | |
import pandas as pd | |
from diffusers import LCMScheduler, AutoencoderTiny, EulerDiscreteScheduler, UNet2DConditionModel | |
from diffusers.models import ImageProjection | |
from patch_sdxl import SDEmb | |
import torch | |
import spaces | |
import random | |
import time | |
import torch | |
from urllib.request import urlopen | |
from PIL import Image | |
import requests | |
from io import BytesIO, StringIO | |
from huggingface_hub import hf_hub_download | |
from safetensors.torch import load_file | |
prompt_list = [p for p in list(set( | |
pd.read_csv('./twitter_prompts.csv').iloc[:, 1].tolist())) if type(p) == str] | |
start_time = time.time() | |
####################### Setup Model | |
model_id = "stabilityai/stable-diffusion-xl-base-1.0" | |
sdxl_lightening = "ByteDance/SDXL-Lightning" | |
ckpt = "sdxl_lightning_2step_unet.safetensors" | |
unet = UNet2DConditionModel.from_config(model_id, subfolder="unet").to("cuda", torch.float16) | |
unet.load_state_dict(load_file(hf_hub_download(sdxl_lightening, ckpt), device="cuda")) | |
pipe = SDEmb.from_pretrained(model_id, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda") | |
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16) | |
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") | |
pipe.to(device='cuda') | |
pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin") | |
output_hidden_state = False | |
####################### | |
def predict( | |
prompt, | |
im_emb=None, | |
progress=gr.Progress(track_tqdm=True) | |
): | |
"""Run a single prediction on the model""" | |
with torch.no_grad(): | |
if im_emb == None: | |
im_emb = torch.zeros(1, 1280, dtype=torch.float16, device='cuda') | |
image = pipe( | |
prompt=prompt, | |
ip_adapter_emb=im_emb.to('cuda'), | |
height=1024, | |
width=1024, | |
num_inference_steps=2, | |
guidance_scale=0, | |
).images[0] | |
im_emb, _ = pipe.encode_image( | |
image, 'cuda', 1, output_hidden_state | |
) | |
return image, im_emb.to(DEVICE) | |
# TODO add to state instead of shared across all | |
glob_idx = 0 | |
def next_image(embs, ys, calibrate_prompts): | |
global glob_idx | |
glob_idx = glob_idx + 1 | |
# handle case where every instance of calibration prompts is 'Neither' or 'Like' or 'Dislike' | |
if len(calibrate_prompts) == 0 and len(list(set(ys))) <= 1: | |
embs.append(.01*torch.randn(1, 1280)) | |
embs.append(.01*torch.randn(1, 1280)) | |
ys.append(0) | |
ys.append(1) | |
with torch.no_grad(): | |
if len(calibrate_prompts) > 0: | |
print('######### Calibrating with sample prompts #########') | |
prompt = calibrate_prompts.pop(0) | |
print(prompt) | |
image, img_emb = predict(prompt) | |
embs.append(img_emb) | |
return image, embs, ys, calibrate_prompts | |
else: | |
print('######### Roaming #########') | |
# sample only as many negatives as there are positives | |
indices = range(len(ys)) | |
pos_indices = [i for i in indices if ys[i] == 1] | |
neg_indices = [i for i in indices if ys[i] == 0] | |
lower = min(len(pos_indices), len(neg_indices)) | |
neg_indices = random.sample(neg_indices, lower) | |
pos_indices = random.sample(pos_indices, lower) | |
cut_embs = [embs[i] for i in neg_indices] + [embs[i] for i in pos_indices] | |
cut_ys = [ys[i] for i in neg_indices] + [ys[i] for i in pos_indices] | |
feature_embs = torch.stack([e[0].detach().cpu() for e in cut_embs]) | |
scaler = preprocessing.StandardScaler().fit(feature_embs) | |
feature_embs = scaler.transform(feature_embs) | |
print(np.array(feature_embs).shape, np.array(ys).shape) | |
lin_class = LinearSVC(max_iter=50000, dual='auto', class_weight='balanced').fit(np.array(feature_embs), np.array(cut_ys)) | |
lin_class.coef_ = torch.tensor(lin_class.coef_, dtype=torch.double) | |
lin_class.coef_ = (lin_class.coef_.flatten() / (lin_class.coef_.flatten().norm())).unsqueeze(0) | |
rng_prompt = random.choice(prompt_list) | |
w = 1# if len(embs) % 2 == 0 else 0 | |
im_emb = w * lin_class.coef_.to(device=DEVICE, dtype=torch.float16) | |
prompt= 'an image' if glob_idx % 2 == 0 else rng_prompt | |
print(prompt) | |
image, im_emb = predict(prompt, im_emb) | |
embs.append(im_emb) | |
return image, embs, ys, calibrate_prompts | |
def start(_, embs, ys, calibrate_prompts): | |
image, embs, ys, calibrate_prompts = next_image(embs, ys, calibrate_prompts) | |
return [ | |
gr.Button(value='Like (L)', interactive=True), | |
gr.Button(value='Neither (Space)', interactive=True), | |
gr.Button(value='Dislike (A)', interactive=True), | |
gr.Button(value='Start', interactive=False), | |
image, | |
embs, | |
ys, | |
calibrate_prompts | |
] | |
def choose(choice, embs, ys, calibrate_prompts): | |
if choice == 'Like': | |
choice = 1 | |
elif choice == 'Neither': | |
_ = embs.pop(-1) | |
img, embs, ys, calibrate_prompts = next_image(embs, ys, calibrate_prompts) | |
return img, embs, ys, calibrate_prompts | |
else: | |
choice = 0 | |
ys.append(choice) | |
img, embs, ys, calibrate_prompts = next_image(embs, ys, calibrate_prompts) | |
return img, embs, ys, calibrate_prompts | |
css = '''.gradio-container{max-width: 700px !important} | |
#description{text-align: center} | |
#description h1{display: block} | |
#description p{margin-top: 0} | |
''' | |
js = ''' | |
<script> | |
document.addEventListener('keydown', function(event) { | |
if (event.key === 'a' || event.key === 'A') { | |
// Trigger click on 'dislike' if 'A' is pressed | |
document.getElementById('dislike').click(); | |
} else if (event.key === ' ' || event.keyCode === 32) { | |
// Trigger click on 'neither' if Spacebar is pressed | |
document.getElementById('neither').click(); | |
} else if (event.key === 'l' || event.key === 'L') { | |
// Trigger click on 'like' if 'L' is pressed | |
document.getElementById('like').click(); | |
} | |
}); | |
</script> | |
''' | |
with gr.Blocks(css=css, head=js) as demo: | |
gr.Markdown('''# Generative Recommenders | |
Explore the latent space without text prompts, based on your preferences. [Learn more on the blog](https://rynmurdock.github.io/posts/2024/3/generative_recomenders/) | |
''', elem_id="description") | |
embs = gr.State([]) | |
ys = gr.State([]) | |
calibrate_prompts = gr.State([ | |
"4k photo", | |
'surrealist art', | |
# 'a psychedelic, fractal view', | |
'a beautiful collage', | |
'abstract art', | |
'an eldritch image', | |
'a sketch', | |
# 'a city full of darkness and graffiti', | |
'', | |
]) | |
with gr.Row(elem_id='output-image'): | |
img = gr.Image(interactive=False, elem_id='output-image',width=700) | |
with gr.Row(equal_height=True): | |
b3 = gr.Button(value='Dislike (A)', interactive=False, elem_id="dislike") | |
b2 = gr.Button(value='Neither (Space)', interactive=False, elem_id="neither") | |
b1 = gr.Button(value='Like (L)', interactive=False, elem_id="like") | |
b1.click( | |
choose, | |
[b1, embs, ys, calibrate_prompts], | |
[img, embs, ys, calibrate_prompts] | |
) | |
b2.click( | |
choose, | |
[b2, embs, ys, calibrate_prompts], | |
[img, embs, ys, calibrate_prompts] | |
) | |
b3.click( | |
choose, | |
[b3, embs, ys, calibrate_prompts], | |
[img, embs, ys, calibrate_prompts] | |
) | |
with gr.Row(): | |
b4 = gr.Button(value='Start') | |
b4.click(start, | |
[b4, embs, ys, calibrate_prompts], | |
[b1, b2, b3, b4, img, embs, ys, calibrate_prompts]) | |
with gr.Row(): | |
html = gr.HTML('''<div style='text-align:center; font-size:20px'>You will calibrate for several prompts and then roam.</ div>''') | |
demo.launch() # Share your demo with just 1 extra parameter π |