Spaces:
Paused
Paused
| 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.float32) | |
| # 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() | |
| 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() | |