add demo

app.py

@@ -0,0 +1,125 @@
+import gradio as gr
+from dataclasses import dataclass
+import PIL
+import PIL.Image
+
+import torch
+import numpy as np
+from gradio_editor3d import Editor3D as g3deditor
+import copy
+from loosecontrol import LooseControlNet
+
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+cn = LooseControlNet()
+cn.pipe = cn.pipe.to(torch_device=device, torch_dtype=torch.float16)
+
+# Need to figure out a better way how to do this per user, making 'cf attention' act like a state per user.
+# For now, we just copy the model. 
+cn_with_cf = copy.deepcopy(cn)
+cn_with_cf.set_cf_attention()
+
+
+@dataclass
+class FixedInputs:
+    prompt: str
+    seed: int
+    depth: PIL.Image.Image
+
+
+negative_prompt = "blurry, text, caption, lowquality, lowresolution, low res, grainy, ugly"
+def depth2image(prompt, seed, depth):
+    seed = int(seed)
+    gen = cn(prompt, control_image=depth, controlnet_conditioning_scale=1.0, generator=torch.Generator().manual_seed(seed), num_inference_steps=20, negative_prompt=negative_prompt)
+    return gen
+
+def edit_previous(prompt, seed, depth, fixed_inputs):
+    seed = int(seed)
+    control_image = [fixed_inputs.depth, depth]
+    prompt = [fixed_inputs.prompt, prompt]
+    neg_prompt = [negative_prompt, negative_prompt]
+    generator = [torch.Generator().manual_seed(fixed_inputs.seed), torch.Generator().manual_seed(seed)]
+    gen = cn_with_cf(prompt, control_image=control_image, controlnet_conditioning_scale=1.0, generator=generator, num_inference_steps=20, negative_prompt=neg_prompt)[-1]
+    return gen
+
+def run(prompt, seed, depth, should_edit, fixed_inputs):
+    depth = depth.convert("RGB")
+    # all values below [3,3,3] in depth should actually be set to [255,255,255]
+    # This is to due the nature of training data and is experimental right now. 
+    # Not in use for now.
+    # depth = np.array(depth)
+    # depth[depth < 3] = 255
+    # depth = PIL.Image.fromarray(depth)
+
+    fixed_inputs = fixed_inputs[0]
+    if should_edit and fixed_inputs is not None:
+        return edit_previous(prompt, seed, depth, fixed_inputs)
+    else:
+        return depth2image(prompt, seed, depth)
+    
+def handle_edit_change(edit, prompt, seed, image_input, fixed_inputs):
+    if edit:
+        fixed_inputs[0] = FixedInputs(prompt, int(seed), image_input)
+    else:
+        fixed_inputs[0] = None
+    return fixed_inputs
+
+
+css = """
+
+#image_output {
+width: 512px;
+height: 512px; 
+"""
+
+
+main_description = """
+# LooseControl
+
+This is the official demo for the paper [LooseControl: Lifting ControlNet for Generalized Depth Conditioning](https://shariqfarooq123.github.io/loose-control/).
+Our 3D Box Editing allows users to interactively edit the 3D boxes representing objects in the scene. Users can change the position, size, and orientation of 3D boxes, allowing to quickly create and edit the scenes to their liking in a 3D-aware manner.
+Best viewed on desktop.
+"""
+
+instructions_editor3d = """
+## Instructions for Editor3D UI
+- Use 'WASD' keys to move the camera.
+- Click on an object to select it.
+- Use the sliders to change the position, size, and orientation of the selected object. Sliders support click and drag for faster editing.
+- Use the 'Add Box', 'Delete', and 'Duplicate' buttons to add, delete, and duplicate objects.
+- Delete and Duplicate buttons work on the selected object. Duplicate creates a copy and selects it.
+- Use the 'Toggle Mode' to switch between "normal" and "depth" mode. Final image sent to the model should be in "depth" mode.
+- Use the 'Render' button to render the scene and send it to the model for generation.
+
+### Lock style checkbox - Fixes the style of the latest generated image. 
+This allows users to edit the 3D boxes without changing the style of the generated image. This is useful when the user is satisfied with the style/content of the generated image and wants to edit the 3D boxes without changing the overall essence of the scene.
+It can be used to create stop motion videos like those shown [here](https://shariqfarooq123.github.io/loose-control/).
+
+"""
+
+
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(main_description)
+
+    fixed_inputs = gr.State([None])
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", placeholder="Write your prompt", elem_id="input")
+        seed = gr.Textbox(value=42, label="Seed", elem_id="seed")
+        should_edit = gr.Checkbox(label="Lock style", elem_id="edit")
+    
+    with gr.Row():
+        image_input = g3deditor(elem_id="image_input")
+    
+    with gr.Row():
+        image_output = gr.Image(elem_id="image_output", type='pil')
+
+    should_edit.change(fn=handle_edit_change, inputs=[should_edit, prompt, seed, image_input, fixed_inputs], outputs=[fixed_inputs])
+    image_input.change(fn=run, inputs=[prompt, seed, image_input, should_edit, fixed_inputs], outputs=[image_output])
+    with gr.Accordion("Instructions"):
+        gr.Markdown(instructions_editor3d)
+
+demo.queue().launch()
+
+
+

cross_frame_attention.py

@@ -0,0 +1,120 @@
+import torch
+from einops import rearrange
+
+class CrossFrameAttnProcessor:
+    def __init__(self, unet_chunk_size=2):
+        self.unet_chunk_size = unet_chunk_size
+
+    def __call__(
+            self,
+            attn,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None, **kwargs):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        query = attn.to_q(hidden_states)
+
+        is_cross_attention = encoder_hidden_states is not None
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        # Sparse Attention
+        if not is_cross_attention:
+            video_length = key.size()[0] // self.unet_chunk_size
+            # print("Video length is", video_length)
+            # former_frame_index = torch.arange(video_length) - 1
+            # former_frame_index[0] = 0
+            former_frame_index = [0] * video_length
+            key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
+            key = key[:, former_frame_index]
+            key = rearrange(key, "b f d c -> (b f) d c")
+            value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
+            value = value[:, former_frame_index]
+            value = rearrange(value, "b f d c -> (b f) d c")
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+    
+
+
+class AttnProcessorX:
+    r"""
+    Default processor for performing attention-related computations.
+    """
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        scale=1.0,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states, scale=scale)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states, scale=scale)
+        value = attn.to_v(encoder_hidden_states, scale=scale)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states, scale=scale)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states

loosecontrol.py

@@ -0,0 +1,135 @@
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler,
+)
+import torch
+import PIL
+import PIL.Image
+from diffusers.loaders import UNet2DConditionLoadersMixin
+from typing import Dict
+from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
+import functools
+from cross_frame_attention import CrossFrameAttnProcessor
+
+TEXT_ENCODER_NAME = "text_encoder"
+UNET_NAME = "unet"
+NEGATIVE_PROMPT = "blurry, text, caption, lowquality, lowresolution, low res, grainy, ugly"
+
+def attach_loaders_mixin(model):
+    # hacky way to make ControlNet work with LoRA. This may not be required in future versions of diffusers.
+    model.text_encoder_name = TEXT_ENCODER_NAME
+    model.unet_name = UNET_NAME
+    r"""
+    Attach the [`UNet2DConditionLoadersMixin`] to a model. This will add the
+    all the methods from the mixin 'UNet2DConditionLoadersMixin' to the model.
+    """
+    # mixin_instance = UNet2DConditionLoadersMixin()
+    for attr_name, attr_value in vars(UNet2DConditionLoadersMixin).items():
+        # print(attr_name)
+        if callable(attr_value):
+            # setattr(model, attr_name, functools.partialmethod(attr_value, model).__get__(model, model.__class__))
+            setattr(model, attr_name, functools.partial(attr_value, model))
+    return model
+
+def set_attn_processor(module, processor, _remove_lora=False):
+    def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+        if hasattr(module, "set_processor"):
+            if not isinstance(processor, dict):
+                module.set_processor(processor, _remove_lora=_remove_lora)
+            else:
+                module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
+
+        for sub_name, child in module.named_children():
+            fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+    for name, module in module.named_children():
+        fn_recursive_attn_processor(name, module, processor)
+
+
+
+class ControlNetX(ControlNetModel, UNet2DConditionLoadersMixin):
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    # This may not be required in future versions of diffusers.
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+class ControlNetPipeline:
+    def __init__(self, checkpoint="lllyasviel/control_v11f1p_sd15_depth", sd_checkpoint="runwayml/stable-diffusion-v1-5") -> None:
+        controlnet = ControlNetX.from_pretrained(checkpoint)
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                        sd_checkpoint, controlnet=controlnet, requires_safety_checker=False, safety_checker=None,
+                        torch_dtype=torch.float16)
+        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+
+    @torch.no_grad()
+    def __call__(self, 
+                    prompt: str="",
+                    height=512,
+                    width=512, 
+                    control_image=None, 
+                    controlnet_conditioning_scale=1.0, 
+                    num_inference_steps: int=20,
+                   **kwargs) -> PIL.Image.Image:
+        
+        out =  self.pipe(prompt, control_image,
+                            height=height, width=width,
+                            num_inference_steps=num_inference_steps,
+                            controlnet_conditioning_scale=controlnet_conditioning_scale,
+                            **kwargs).images
+
+        return out[0] if len(out) == 1 else out
+    
+    def to(self, *args, **kwargs):
+        self.pipe.to(*args, **kwargs)
+        return self
+
+
+class LooseControlNet(ControlNetPipeline):
+    def __init__(self, loose_control_weights="shariqfarooq/loose-control-3dbox", cn_checkpoint="lllyasviel/control_v11f1p_sd15_depth", sd_checkpoint="runwayml/stable-diffusion-v1-5") -> None:
+        super().__init__(cn_checkpoint, sd_checkpoint)
+        self.pipe.controlnet = attach_loaders_mixin(self.pipe.controlnet)
+        self.pipe.controlnet.load_attn_procs(loose_control_weights)
+
+    def set_normal_attention(self):
+        self.pipe.unet.set_attn_processor(AttnProcessor())
+
+    def set_cf_attention(self, _remove_lora=False):
+        for upblocks in self.pipe.unet.up_blocks[-2:]:
+            set_attn_processor(upblocks, CrossFrameAttnProcessor(), _remove_lora=_remove_lora)
+
+    def edit(self, depth, depth_edit, prompt, prompt_edit=None, seed=42, seed_edit=None, negative_prompt=NEGATIVE_PROMPT, controlnet_conditioning_scale=1.0, num_inference_steps=20, **kwargs):
+        if prompt_edit is None:
+            prompt_edit = prompt
+
+        if seed_edit is None:
+            seed_edit = seed
+    
+        seed = int(seed)
+        seed_edit = int(seed_edit)
+        control_image = [depth, depth_edit]
+        prompt = [prompt, prompt_edit]
+        generator = [torch.Generator().manual_seed(seed), torch.Generator().manual_seed(seed_edit)]
+        gen = self.pipe(prompt, control_image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, generator=generator, num_inference_steps=num_inference_steps, negative_prompt=negative_prompt, **kwargs)[-1]
+        return gen