fix inversion

gradio_app.py

@@ -13,10 +13,12 @@ from main import LPMConfig, main, setup
 DESCRIPTION = '''# Localizing Object-level Shape Variations with Text-to-Image Diffusion Models
 This is a demo for our ''Localizing Object-level Shape Variations with Text-to-Image Diffusion Models'' [paper](https://arxiv.org/abs/2303.11306).
 We introduce a method that generates object-level shape variation for a given image. 
-This demo allows using a real image as well as a generated image. For a real image, a matching prompt is required.
+This demo supports both generated images and real images. To modify a real image, please upload it to the input image block and provide a prompt that describes its contents.
+
 '''
 
 stable, stable_config = setup(LPMConfig())
+stable_for_inversion, _ = setup(LPMConfig())
 
 def main_pipeline(
         prompt: str,
@@ -30,8 +32,6 @@ def main_pipeline(
         seed: int,
         input_image: str):
         prompt = prompt.replace(object_of_interest, '{word}')
-        print(number_of_variations)
-        print(proxy_words)
         proxy_words = proxy_words.split(',') if proxy_words != '' else []
         objects_to_preserve = objects_to_preserve.split(',') if objects_to_preserve != '' else []
         background_nouns = background_nouns.split(',') if background_nouns != '' else []
@@ -48,7 +48,7 @@ def main_pipeline(
             real_image_path="" if input_image is None else input_image
         )
 
-        result_images, result_proxy_words = main(stable, stable_config, args)
+        result_images, result_proxy_words = main(stable, stable_config, stable_for_inversion, args)
         result_images = [im.permute(1, 2, 0).cpu().numpy() for im in result_images]
         result_images = [(im * 255).astype(np.uint8) for im in result_images]
         result_images = [Image.fromarray(im) for im in result_images]
@@ -60,8 +60,8 @@ with gr.Blocks(css='style.css') as demo:
     gr.Markdown(DESCRIPTION)
 
     gr.HTML(
-        '''<center><a href="https://huggingface.co/spaces/orpatashnik/local-prompt-mixing?duplicate=true">
-        <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space and run securely with your OpenAI API Key</center>''')
+        '''<a href="https://huggingface.co/spaces/orpatashnik/local-prompt-mixing?duplicate=true">
+        <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space and run securely with your OpenAI API Key''')
 
     with gr.Row():
         with gr.Column():

main.py

@@ -34,7 +34,7 @@ def setup(args):
     return ldm_stable, ldm_stable_config
 
 
-def main(ldm_stable, ldm_stable_config, args):
+def main(ldm_stable, ldm_stable_config, ldm_stable_inversion, args):
 
     similar_words, prompts, another_prompts = get_proxy_prompts(args, ldm_stable)
     exp_path = save_args_dict(args, similar_words)
@@ -44,7 +44,7 @@ def main(ldm_stable, ldm_stable_config, args):
     uncond_embeddings = None
 
     if args.real_image_path != "":
-        x_t, uncond_embeddings = invert_image(args, ldm_stable, ldm_stable_config, prompts, exp_path)
+        x_t, uncond_embeddings = invert_image(args, ldm_stable_inversion, ldm_stable_config, prompts, exp_path)
 
     image, x_t, orig_all_latents, orig_mask, average_attention = generate_original_image(args, ldm_stable, ldm_stable_config, prompts, x_t, uncond_embeddings)
     save_image(ToTensor()(image[0]), f"{exp_path}/{similar_words[0]}.jpg")