|
|
|
"""Copy of compose_glide.ipynb |
|
|
|
Automatically generated by Colaboratory. |
|
|
|
Original file is located at |
|
https://colab.research.google.com/drive/19xx6Nu4FeiGj-TzTUFxBf-15IkeuFx_F |
|
""" |
|
|
|
import gradio as gr |
|
import torch as th |
|
|
|
from composable_diffusion.download import download_model |
|
from composable_diffusion.model_creation import create_model_and_diffusion as create_model_and_diffusion_for_clevr |
|
from composable_diffusion.model_creation import model_and_diffusion_defaults as model_and_diffusion_defaults_for_clevr |
|
from composable_diffusion.composable_stable_diffusion.pipeline_composable_stable_diffusion import \ |
|
ComposableStableDiffusionPipeline |
|
|
|
import os |
|
import shutil |
|
import time |
|
import glob |
|
import numpy as np |
|
import open3d as o3d |
|
import open3d.visualization.rendering as rendering |
|
|
|
import plotly.graph_objects as go |
|
from PIL import Image |
|
from tqdm.auto import tqdm |
|
from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config |
|
from point_e.diffusion.sampler import PointCloudSampler |
|
from point_e.models.download import load_checkpoint |
|
from point_e.models.configs import MODEL_CONFIGS, model_from_config |
|
from point_e.util.pc_to_mesh import marching_cubes_mesh |
|
|
|
has_cuda = th.cuda.is_available() |
|
device = th.device('cpu' if not th.cuda.is_available() else 'cuda') |
|
print(has_cuda) |
|
|
|
|
|
pipe = ComposableStableDiffusionPipeline.from_pretrained( |
|
"CompVis/stable-diffusion-v1-4", |
|
).to(device) |
|
|
|
|
|
|
|
|
|
|
|
clevr_options = model_and_diffusion_defaults_for_clevr() |
|
|
|
flags = { |
|
"image_size": 128, |
|
"num_channels": 192, |
|
"num_res_blocks": 2, |
|
"learn_sigma": True, |
|
"use_scale_shift_norm": False, |
|
"raw_unet": True, |
|
"noise_schedule": "squaredcos_cap_v2", |
|
"rescale_learned_sigmas": False, |
|
"rescale_timesteps": False, |
|
"num_classes": '2', |
|
"dataset": "clevr_pos", |
|
"use_fp16": has_cuda, |
|
"timestep_respacing": '100' |
|
} |
|
|
|
for key, val in flags.items(): |
|
clevr_options[key] = val |
|
|
|
clevr_model, clevr_diffusion = create_model_and_diffusion_for_clevr(**clevr_options) |
|
clevr_model.eval() |
|
if has_cuda: |
|
clevr_model.convert_to_fp16() |
|
|
|
clevr_model.to(device) |
|
clevr_model.load_state_dict(th.load(download_model('clevr_pos'), device)) |
|
device = th.device('cpu' if not th.cuda.is_available() else 'cuda') |
|
|
|
print('creating base model...') |
|
base_name = 'base40M-textvec' |
|
base_model = model_from_config(MODEL_CONFIGS[base_name], device) |
|
base_model.eval() |
|
base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name]) |
|
|
|
print('creating upsample model...') |
|
upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device) |
|
upsampler_model.eval() |
|
upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample']) |
|
|
|
print('downloading base checkpoint...') |
|
base_model.load_state_dict(load_checkpoint(base_name, device)) |
|
|
|
print('downloading upsampler checkpoint...') |
|
upsampler_model.load_state_dict(load_checkpoint('upsample', device)) |
|
|
|
print('creating SDF model...') |
|
name = 'sdf' |
|
model = model_from_config(MODEL_CONFIGS[name], device) |
|
model.eval() |
|
|
|
print('loading SDF model...') |
|
model.load_state_dict(load_checkpoint(name, device)) |
|
|
|
|
|
def compose_pointe(prompt, weights, version): |
|
weight_list = [float(x.strip()) for x in weights.split('|')] |
|
sampler = PointCloudSampler( |
|
device=device, |
|
models=[base_model, upsampler_model], |
|
diffusions=[base_diffusion, upsampler_diffusion], |
|
num_points=[1024, 4096 - 1024], |
|
aux_channels=['R', 'G', 'B'], |
|
guidance_scale=[weight_list, 0.0], |
|
model_kwargs_key_filter=('texts', ''), |
|
) |
|
|
|
def generate_pcd(prompt_list): |
|
|
|
samples = None |
|
for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=prompt_list))): |
|
samples = x |
|
return samples |
|
|
|
def generate_fig(samples): |
|
pc = sampler.output_to_point_clouds(samples)[0] |
|
return pc |
|
|
|
def generate_mesh(pc): |
|
mesh = marching_cubes_mesh( |
|
pc=pc, |
|
model=model, |
|
batch_size=4096, |
|
grid_size=128, |
|
progress=True, |
|
) |
|
return mesh |
|
|
|
def generate_video(mesh_path): |
|
render = rendering.OffscreenRenderer(640, 480) |
|
mesh = o3d.io.read_triangle_mesh(mesh_path) |
|
mesh.compute_vertex_normals() |
|
|
|
mat = o3d.visualization.rendering.MaterialRecord() |
|
mat.shader = 'defaultLit' |
|
|
|
render.scene.camera.look_at([0, 0, 0], [1, 1, 1], [0, 0, 1]) |
|
render.scene.add_geometry('mesh', mesh, mat) |
|
|
|
timestr = time.strftime("%Y%m%d-%H%M%S") |
|
os.makedirs(timestr, exist_ok=True) |
|
|
|
def update_geometry(): |
|
render.scene.clear_geometry() |
|
render.scene.add_geometry('mesh', mesh, mat) |
|
|
|
def generate_images(): |
|
for i in range(64): |
|
|
|
R = mesh.get_rotation_matrix_from_xyz((0, 0, np.pi / 32)) |
|
mesh.rotate(R, center=(0, 0, 0)) |
|
|
|
update_geometry() |
|
img = render.render_to_image() |
|
o3d.io.write_image(os.path.join(timestr + "/{:05d}.jpg".format(i)), img, quality=100) |
|
time.sleep(0.05) |
|
|
|
generate_images() |
|
image_list = [] |
|
for filename in sorted(glob.glob(f'{timestr}/*.jpg')): |
|
im = Image.open(filename) |
|
image_list.append(im) |
|
|
|
shutil.rmtree(timestr) |
|
return image_list |
|
|
|
prompt_list = [x.strip() for x in prompt.split("|")] |
|
pcd = generate_pcd(prompt_list) |
|
pc = generate_fig(pcd) |
|
|
|
fig = go.Figure( |
|
data=[ |
|
go.Scatter3d( |
|
x=pc.coords[:, 0], y=pc.coords[:, 1], z=pc.coords[:, 2], |
|
mode='markers', |
|
marker=dict( |
|
size=2, |
|
color=['rgb({},{},{})'.format(r, g, b) for r, g, b in |
|
zip(pc.channels["R"], pc.channels["G"], pc.channels["B"])], |
|
) |
|
) |
|
], |
|
layout=dict( |
|
scene=dict( |
|
xaxis=dict(visible=False), |
|
yaxis=dict(visible=False), |
|
zaxis=dict(visible=False) |
|
) |
|
), |
|
) |
|
return fig |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def compose_clevr_objects(prompt, weights, steps): |
|
weights = [float(x.strip()) for x in weights.split('|')] |
|
weights = th.tensor(weights, device=device).reshape(-1, 1, 1, 1) |
|
coordinates = [ |
|
[ |
|
float(x.split(',')[0].strip()), float(x.split(',')[1].strip())] |
|
for x in prompt.split('|') |
|
] |
|
coordinates += [[-1, -1]] |
|
batch_size = 1 |
|
|
|
clevr_options['timestep_respacing'] = str(int(steps)) |
|
_, clevr_diffusion = create_model_and_diffusion_for_clevr(**clevr_options) |
|
|
|
def model_fn(x_t, ts, **kwargs): |
|
half = x_t[:1] |
|
combined = th.cat([half] * kwargs['y'].size(0), dim=0) |
|
model_out = clevr_model(combined, ts, **kwargs) |
|
eps, rest = model_out[:, :3], model_out[:, 3:] |
|
masks = kwargs.get('masks') |
|
cond_eps = eps[masks] |
|
uncond_eps = eps[~masks] |
|
half_eps = uncond_eps + (weights * (cond_eps - uncond_eps)).sum(dim=0, keepdims=True) |
|
eps = th.cat([half_eps] * x_t.size(0), dim=0) |
|
return th.cat([eps, rest], dim=1) |
|
|
|
def sample(coordinates): |
|
masks = [True] * (len(coordinates) - 1) + [False] |
|
model_kwargs = dict( |
|
y=th.tensor(coordinates, dtype=th.float, device=device), |
|
masks=th.tensor(masks, dtype=th.bool, device=device) |
|
) |
|
samples = clevr_diffusion.p_sample_loop( |
|
model_fn, |
|
(len(coordinates), 3, clevr_options["image_size"], clevr_options["image_size"]), |
|
device=device, |
|
clip_denoised=True, |
|
progress=True, |
|
model_kwargs=model_kwargs, |
|
cond_fn=None, |
|
)[:batch_size] |
|
|
|
return samples |
|
|
|
samples = sample(coordinates) |
|
out_img = samples[0].permute(1, 2, 0) |
|
out_img = (out_img + 1) / 2 |
|
out_img = (out_img.detach().cpu() * 255.).to(th.uint8) |
|
out_img = out_img.numpy() |
|
|
|
return out_img |
|
|
|
|
|
def stable_diffusion_compose(prompt, steps, weights, seed): |
|
generator = th.Generator("cuda").manual_seed(int(seed)) |
|
image = pipe(prompt, guidance_scale=7.5, num_inference_steps=steps, |
|
weights=weights, generator=generator).images[0] |
|
image.save(f'{"_".join(prompt.split())}.png') |
|
return image |
|
|
|
|
|
def compose_2D_diffusion(prompt, weights, version, steps, seed): |
|
try: |
|
with th.no_grad(): |
|
if version == 'Stable_Diffusion_1v_4': |
|
res = stable_diffusion_compose(prompt, steps, weights, seed) |
|
return res |
|
else: |
|
return compose_clevr_objects(prompt, weights, steps) |
|
except Exception as e: |
|
return None |
|
|
|
|
|
examples_1 = "A castle in a forest | grainy, fog" |
|
examples_3 = '0.1, 0.5 | 0.3, 0.5 | 0.5, 0.5 | 0.7, 0.5 | 0.9, 0.5' |
|
examples_5 = 'a white church | lightning in the background' |
|
examples_6 = 'mystical trees | A dark magical pond | dark' |
|
examples_7 = 'A lake | A mountain | Cherry Blossoms next to the lake' |
|
|
|
image_examples = [ |
|
[examples_6, "7.5 | 7.5 | -7.5", 'Stable_Diffusion_1v_4', 50, 8], |
|
[examples_6, "7.5 | 7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 8], |
|
[examples_1, "7.5 | -7.5", 'Stable_Diffusion_1v_4', 50, 0], |
|
[examples_7, "7.5 | 7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 3], |
|
[examples_5, "7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 0], |
|
[examples_3, "7.5 | 7.5 | 7.5 | 7.5 | 7.5", 'CLEVR Objects', 100, 0] |
|
] |
|
|
|
pointe_examples = [["a cake | a house", "7.5 | 7.5", 'Point-E'], |
|
["a chair | chair legs", "7.5 | -7.5", 'Point-E'], |
|
["a green avocado | a chair", "7.5 | 3", 'Point-E'], |
|
["a toilet | a chair", "7 | 5", 'Point-E']] |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown( |
|
"""<h1 style="text-align: center;"><b>Composable Diffusion Models (ECCV |
|
2022)</b> - <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion |
|
-Models/">Project Page</a></h1>""") |
|
gr.Markdown( |
|
"""<table style="display: inline-table; table-layout: fixed; width: 100%;"> |
|
<tr> |
|
<td> |
|
<figure> |
|
<img src="https://media.giphy.com/media/gKfDjdXy0lbYNyROKo/giphy.gif" style="text-align:center; width:100%; display:block; margin:auto;"> |
|
<figcaption style="color: black; font-size: 15px; text-align: center;">"Mystical trees" <span style="color: red">AND</span> "A magical pond" <span style="color: red">AND</span> "Dark"</figcaption> |
|
</figure> |
|
</td> |
|
<td> |
|
<figure> |
|
<img src="https://media.giphy.com/media/sf5m1Z5FldemLMatWn/giphy.gif" style="text-align:center; width:100%; display:block; margin:auto;"> |
|
<figcaption style="color: black; font-size: 15px; text-align: center;">"Mystical trees" <span style="color: red">AND</span> "A magical pond" <span style="color: red">AND NOT</span> "Dark"</figcaption> |
|
</figure> |
|
</td> |
|
<td> |
|
<figure> |
|
<img src="https://media.giphy.com/media/LDmNSM9NmNpaljMKiF/giphy.gif" style="text-align:center; width:100%; display:block; margin:auto;"> |
|
<figcaption style="color: black; font-size: 15px; text-align: center;">"A chair" <span style="color: red">AND NOT</span> "Chair legs"</figcaption> |
|
</figure> |
|
</td> |
|
<td> |
|
<figure> |
|
<img src="https://media.giphy.com/media/nFkMh70kzZCwjbRrx5/giphy.gif" style="text-align:center; width:100%; display:block; margin:auto;"> |
|
<figcaption style="color: black; font-size: 15px; text-align: center;">"A monitor" <span style="color: red">AND</span> "A brown couch"</figcaption> |
|
</figure> |
|
</td> |
|
</tr> |
|
</table> |
|
""" |
|
) |
|
gr.Markdown( |
|
"""<p style="font-size: 18px;">Compositional visual generation by composing pre-trained diffusion models |
|
using compositional operators, <b>AND</b> and <b>NOT</b>.</p>""") |
|
gr.Markdown( |
|
"""<p style="font-size: 18px;">When composing multiple inputs, please use <b>β|β</b> to separate them </p>""") |
|
gr.Markdown( |
|
"""<p>( <b>Clevr Note</b>: For composing CLEVR objects, we recommend using <b><i>x</i></b> in range <b><i>[0.1, |
|
0.9]</i></b> and <b><i>y</i></b> in range <b><i>[0.25, 0.7]</i></b>, since the training dataset labels are in |
|
given ranges.)</p>""") |
|
gr.Markdown( |
|
"""<p>( <b>Point-E Note</b>: This demo only shows the point cloud results instead of meshes due to |
|
hardware limitation. For mesh results, check out our code to render them on your local machine!)</p>""") |
|
gr.Markdown( |
|
"""<p>( <b>Stable Diffusion Note</b>: Stable Diffusion has a filter enabled, so it sometimes generates all black |
|
results for possibly inappropriate images.)</p>""") |
|
gr.Markdown( |
|
"""<p>( <b>Note</b>: Absolute values of weights should be > 1, negative weights indicate negation.)</p>""" |
|
) |
|
with gr.Row(): |
|
with gr.Column(): |
|
gr.Markdown( |
|
"""<h4>Composing natural language descriptions / objects for 2D image |
|
generation</h4>""") |
|
with gr.Row(): |
|
text_input = gr.Textbox(value="mystical trees | A dark magical pond | dark", label="Text to image prompt") |
|
weights_input = gr.Textbox(value="7.5 | 7.5 | 7.5", label="Weights") |
|
with gr.Row(): |
|
seed_input = gr.Number(0, label="Seed") |
|
steps_input = gr.Slider(10, 200, value=50, label="Steps") |
|
with gr.Row(): |
|
model_input = gr.Radio( |
|
['Stable_Diffusion_1v_4', 'CLEVR Objects'], type="value", label='Text to image model', |
|
value='Stable_Diffusion_1v_4') |
|
image_output = gr.Image() |
|
image_button = gr.Button("Generate") |
|
img_examples = gr.Examples( |
|
examples=image_examples, |
|
inputs=[text_input, weights_input, model_input, steps_input, seed_input] |
|
) |
|
|
|
with gr.Column(): |
|
gr.Markdown( |
|
"""<h4>Composing natural language descriptions for 3D asset generation</h4>""") |
|
with gr.Row(): |
|
asset_input = gr.Textbox(value="a cake | a house", label="Text to 3D prompt") |
|
with gr.Row(): |
|
asset_weights = gr.Textbox(value="7.5 | 7.5", label="Weights") |
|
with gr.Row(): |
|
asset_model = gr.Radio(['Point-E'], type="value", label='Text to 3D model', value='Point-E') |
|
|
|
asset_output = gr.Plot(label='Plot') |
|
asset_button = gr.Button("Generate") |
|
asset_examples = gr.Examples(examples=pointe_examples, inputs=[asset_input, asset_weights, asset_model]) |
|
|
|
image_button.click(compose_2D_diffusion, |
|
inputs=[text_input, weights_input, model_input, steps_input, seed_input], |
|
outputs=image_output) |
|
asset_button.click(compose_pointe, inputs=[asset_input, asset_weights, asset_model], outputs=asset_output) |
|
|
|
if __name__ == "__main__": |
|
demo.queue(max_size=5) |
|
demo.launch(debug=True) |
|
|