technicalkiwi
/
latentblending
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more powerful UI

This commit is contained in:
DGX 2024-03-27 21:23:21 +00:00
parent ac56d0e2c0
commit 2a2886157f
1 changed files with 134 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

195
latentblending/gradio_ui.py
View File

@ -13,12 +13,9 @@ import uuid
from diffusers import AutoPipelineForText2Image
from latentblending.blending_engine import BlendingEngine
import datetime
warnings.filterwarnings('ignore')
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = False
import tempfile
import json
from lunar_tools import concatenate_movies
class BlendingFrontend():
@ -43,47 +40,26 @@ class BlendingFrontend():
self.prompt1 = ""
self.prompt2 = ""
self.negative_prompt = ""
self.nmb_preview_images = 4
# Vars
self.prompt = None
self.negative_prompt = None
self.list_seeds = []
self.idx_movie = 0
self.list_seeds = []
self.list_images_preview = []
self.data = []
def take_image0(self):
return self.take_image(0)
def take_image1(self):
return self.take_image(1)
def take_image2(self):
return self.take_image(2)
def take_image3(self):
return self.take_image(3)
self.idx_img_selected = None
self.jpg_quality = 80
self.fp_movie = ''
self.duration_single_trans = 10
def take_image(self, id_img):
if self.prompt is None:
print("Cannot take because no prompt was set!")
return [None, None, None, None, ""]
if self.idx_movie == 0:
current_time = datetime.datetime.now()
self.fp_out = "movie_" + current_time.strftime("%y%m%d_%H%M") + ".json"
self.data.append({"settings": "sdxl", "width": bf.be.dh.width_img, "height": self.be.dh.height_img, "num_inference_steps": self.be.dh.num_inference_steps})
seed = self.list_seeds[id_img]
self.data.append({"iteration": self.idx_movie, "seed": seed, "prompt": self.prompt, "negative_prompt": self.negative_prompt})
# Write the data list to a JSON file
with open(self.fp_out, 'w') as f:
json.dump(self.data, f, indent=4)
self.idx_movie += 1
self.prompt = None
return [None, None, None, None, ""]
def preview_img_selected(self, data: gr.SelectData, button):
self.idx_img_selected = data.index
print(f"gallery image {self.idx_img_selected} selected, seed {self.list_seeds[self.idx_img_selected]}")
return gr.Button(interactive=True)
def compute_imgs(self, prompt, negative_prompt):
@ -93,22 +69,102 @@ class BlendingFrontend():
self.be.set_prompt2(prompt)
self.be.set_negative_prompt(negative_prompt)
self.list_seeds = []
self.list_images = []
for i in range(4):
seed = np.random.randint(0, 1000000000)
self.list_images_preview = []
self.idx_img_selected = None
for i in range(self.nmb_preview_images):
seed = np.random.randint(0, np.iinfo(np.int32).max)
self.be.seed1 = seed
self.list_seeds.append(seed)
img = self.be.compute_latents1(return_image=True)
self.list_images.append(img)
return self.list_images
fn_img_tmp = f"image_{uuid.uuid4()}.jpg"
temp_img_path = os.path.join(tempfile.gettempdir(), fn_img_tmp)
img.save(temp_img_path)
img.save(temp_img_path, quality=self.jpg_quality, optimize=True)
self.list_images_preview.append(temp_img_path)
return self.list_images_preview
def get_list_images_movie(self):
return [entry["preview_image"] for entry in self.data[1:]]
def init_new_movie(self):
current_time = datetime.datetime.now()
self.fp_movie = "movie_" + current_time.strftime("%y%m%d_%H%M") + ".mp4"
self.fp_json = "movie_" + current_time.strftime("%y%m%d_%H%M") + ".json"
self.data.append({"settings": "sdxl", "width": bf.be.dh.width_img, "height": self.be.dh.height_img, "num_inference_steps": self.be.dh.num_inference_steps})
def add_image_to_video(self):
if self.prompt is None:
print("Cannot take because no prompt was set!")
return self.get_list_images_movie()
if self.idx_movie == 0:
self.init_new_movie()
self.data.append({"iteration": self.idx_movie,
"seed": self.list_seeds[self.idx_img_selected],
"prompt": self.prompt,
"negative_prompt": self.negative_prompt,
"preview_image": self.list_images_preview[self.idx_img_selected]
})
# Write the data list to a JSON file
with open(self.fp_json, 'w') as f:
json.dump(self.data, f, indent=4)
self.idx_movie += 1
self.prompt = None
return self.get_list_images_movie()
def generate_movie(self):
print("starting movie gen")
list_prompts = []
list_negative_prompts = []
list_seeds = []
# Extract prompts, negative prompts, and seeds from the data
for item in self.data[1:]: # Skip the first item as it contains settings
list_prompts.append(item["prompt"])
list_negative_prompts.append(item["negative_prompt"])
list_seeds.append(item["seed"])
list_movie_parts = []
for i in range(len(list_prompts) - 1):
# For a multi transition we can save some computation time and recycle the latents
if i == 0:
self.be.set_prompt1(list_prompts[i])
self.be.set_negative_prompt(list_negative_prompts[i])
self.be.set_prompt2(list_prompts[i + 1])
recycle_img1 = False
else:
self.be.swap_forward()
self.be.set_negative_prompt(list_negative_prompts[i+1])
self.be.set_prompt2(list_prompts[i + 1])
recycle_img1 = True
fp_movie_part = f"tmp_part_{str(i).zfill(3)}.mp4"
fixed_seeds = list_seeds[i:i + 2]
# Run latent blending
self.be.run_transition(
recycle_img1=recycle_img1,
fixed_seeds=fixed_seeds)
# Save movie
self.be.write_movie_transition(fp_movie_part, self.duration_single_trans)
list_movie_parts.append(fp_movie_part)
# Finally, concatenate the result
concatenate_movies(self.fp_movie, list_movie_parts)
print(f"DONE! MOVIE SAVED IN {self.fp_movie}")
return self.fp_movie
if __name__ == "__main__":
width = 786
height = 1024
width = 512
height = 512
num_inference_steps = 4
pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
@ -126,28 +182,41 @@ if __name__ == "__main__":
with gr.Row():
prompt = gr.Textbox(label="prompt")
negative_prompt = gr.Textbox(label="negative prompt")
b_compute = gr.Button('generate preview images', variant='primary')
# with gr.Row():
with gr.Row():
b_compute = gr.Button('compute new images', variant='primary')
gallery_preview = gr.Gallery(
label="Generated images", show_label=False, elem_id="gallery"
, columns=[bf.nmb_preview_images], rows=[1], object_fit="contain", height="auto", allow_preview=False, interactive=False)
with gr.Row():
with gr.Column():
img0 = gr.Image(label="seed1")
b_take0 = gr.Button('take', variant='primary')
with gr.Column():
img1 = gr.Image(label="seed2")
b_take1 = gr.Button('take', variant='primary')
with gr.Column():
img2 = gr.Image(label="seed3")
b_take2 = gr.Button('take', variant='primary')
with gr.Column():
img3 = gr.Image(label="seed4")
b_take3 = gr.Button('take', variant='primary')
b_select = gr.Button('add selected image to video', variant='primary', interactive=False)
with gr.Row():
gr.Markdown("Your movie contains so far the below frames")
with gr.Row():
gallery_movie = gr.Gallery(
label="Generated images", show_label=False, elem_id="gallery"
, columns=[20], rows=[1], object_fit="contain", height="auto", allow_preview=False, interactive=False)
with gr.Row():
b_generate_movie = gr.Button('generate movie', variant='primary')
with gr.Row():
movie = gr.Video()
# bindings
b_compute.click(bf.compute_imgs, inputs=[prompt, negative_prompt], outputs=gallery_preview)
b_select.click(bf.add_image_to_video, None, gallery_movie)
b_generate_movie.click(bf.generate_movie, None, movie)
gallery_preview.select(bf.preview_img_selected, None, b_select)
b_compute.click(bf.compute_imgs, inputs=[prompt, negative_prompt], outputs=[img0, img1, img2, img3])
b_take0.click(bf.take_image0, outputs=[img0, img1, img2, img3, prompt])
b_take1.click(bf.take_image1, outputs=[img0, img1, img2, img3, prompt])
b_take2.click(bf.take_image2, outputs=[img0, img1, img2, img3, prompt])
b_take3.click(bf.take_image3, outputs=[img0, img1, img2, img3, prompt])
demo.launch(share=bf.share, inbrowser=True, inline=False)