Update app.py

app.py

@@ -7,21 +7,72 @@ Original file is located at
     https://colab.research.google.com/drive/19xx6Nu4FeiGj-TzTUFxBf-15IkeuFx_F
 """
 
-import streamlit as st
 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
+from composable_diffusion.composable_stable_diffusion.pipeline_composable_stable_diffusion import \
+    ComposableStableDiffusionPipeline
 
-# This notebook supports both CPU and GPU.
-# On CPU, generating one sample may take on the order of 20 minutes.
-# On a GPU, it should be under a minute.
+import os
+import shutil
+import time
+import glob
+import numpy as np
+import open3d as o3d
+import open3d.visualization.rendering as rendering
+
+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)
+
+# init stable diffusion model
+pipe = ComposableStableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+).to(device)
+
+pipe.safety_checker = None
+
+# create model for CLEVR Objects
+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')
 
 # init stable diffusion model
 pipe = ComposableStableDiffusionPipeline.from_pretrained(
@@ -61,6 +112,245 @@ clevr_model.to(device)
 clevr_model.load_state_dict(th.load(download_model('clevr_pos'), device))
 print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))
 
+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):
+    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', ''),  # Do not condition the upsampler at all
+    )
+
+    def generate_pcd(prompt_list):
+        # Produce a sample from the model.
+        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
+
+
+# has_cuda = th.cuda.is_available()
+device = th.device('cpu' if not th.cuda.is_available() else 'cuda')
+
+# init stable diffusion model
+pipe = ComposableStableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+).to(device)
+
+pipe.safety_checker = None
+
+# create model for CLEVR Objects
+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')
+
+# init stable diffusion model
+pipe = ComposableStableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+).to(device)
+
+pipe.safety_checker = None
+
+# create model for CLEVR Objects
+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))
+print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))
+
+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', ''),  # Do not condition the upsampler at all
+    )
+
+    def generate_pcd(prompt_list):
+        # Produce a sample from the model.
+        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,  # increase to 128 for resolution used in evals
+            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):
+                # Rotation
+                R = mesh.get_rotation_matrix_from_xyz((0, 0, np.pi / 32))
+                mesh.rotate(R, center=(0, 0, 0))
+                # Update geometry
+                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')):  # assuming gif
+            im = Image.open(filename)
+            image_list.append(im)
+        # remove the folder
+        shutil.rmtree(timestr)
+        return image_list
+
+    prompt_list = [x.strip() for x in prompt.split("|")]
+    pcd = generate_pcd(prompt_list)
+    pc = generate_fig(pcd)
+    mesh = generate_mesh(pc)
+    timestr = time.strftime("%Y%m%d-%H%M%S")
+    mesh_path = os.path.join(f'{timestr}.ply')
+    with open(mesh_path, 'wb') as f:
+        mesh.write_ply(f)
+    image_frames = generate_video(mesh_path)
+    gif_path = os.path.join(f'{timestr}.gif')
+    image_frames[0].save(gif_path, save_all=True, optimizer=False, duration=5, append_images=image_frames[1:], loop=0)
+    return f'{timestr}.gif'
+
 
 def compose_clevr_objects(prompt, weights, steps):
     weights = [float(x.strip()) for x in weights.split('|')]
@@ -123,7 +413,7 @@ def stable_diffusion_compose(prompt, steps, weights, seed):
     return image
 
 
-def compose(prompt, weights, version, steps, seed):
+def compose_2D_diffusion(prompt, weights, version, steps, seed):
     try:
         with th.no_grad():
             if version == 'Stable_Diffusion_1v_4':
@@ -132,35 +422,113 @@ def compose(prompt, weights, version, steps, seed):
             else:
                 return compose_clevr_objects(prompt, weights, steps)
     except Exception as e:
-        print(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'
-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]
+
+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]
 ]
 
-title = 'Compositional Visual Generation with Composable Diffusion Models'
-description = '<p>Our conjunction and negation (a.k.a. negative prompts) operators are also added into stable diffusion webui! (<a href="https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Negative-prompt">Negation</a> and <a href="https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/c26732fbee2a57e621ac22bf70decf7496daa4cd">Conjunction</a>)</p></p><p>See more information from our <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/">Project Page</a>.</p><ul><li>One version is based on the released <a href="https://github.com/openai/glide-text2im">GLIDE</a> and <a href="https://github.com/CompVis/stable-diffusion/">Stable Diffusion</a> for composing natural language description.</li><li>Another is based on our pre-trained CLEVR Object Model for composing objects. <br>(<b>Note</b>: 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.)</li></ul><p>When composing  multiple sentences, use `|` as the delimiter, see given examples below.</p><p>You can also specify the weight of each text by using `|` as the delimiter. When the weight is negative, it will use Negation Operator (NOT), which indicates the corresponding prompt is a negative prompt. Otherwise it will use Conjunction operator (AND).</p><p><b>Only Conjunction operator is enabled for CLEVR Object.</b></p><p><b>Note: When using Stable Diffusion, black images will be returned if the given prompt is detected as problematic. For composing GLIDE model, we recommend using the Colab demo in our <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/">Project Page</a>.</b></p>'
-
-iface = gr.Interface(compose,
-                     inputs=[
-                         gr.Textbox(label='prompt', value='mystical trees | A dark magical pond | dark'),
-                         gr.Textbox(label='weights', value='7.5 | 7.5 | -7.5'),
-                         gr.Radio(['Stable_Diffusion_1v_4', 'CLEVR Objects'], type="value", label='version', value='Stable_Diffusion_1v_4'),
-                         gr.Slider(10, 200, value=50),
-                         gr.Number(8)
-                     ],
-                     outputs='image', cache_examples=False,
-                     title=title, description=description, examples=examples)
-
-iface.launch()
+pointe_examples = [["a cake | a house", "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/lTzdW41bFnrD8AYa0K/giphy.gif" style="text-align:center; width:100%; display:block; margin:auto;">
+                        <figcaption style="color: black; font-size: 15px; text-align: center;">"A toilet" <span style="color: red">AND</span> "A chair"</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>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><hr>""")
+    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.Image(label='GIF')
+            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)