Update app to illusion generation

app.py

@@ -2,6 +2,7 @@ from diffusers import DiffusionPipeline
 from diffusers.utils import pt_to_pil
 import gradio as gr
 import torch
+import numpy as np
 
 
 stage_1 = DiffusionPipeline.from_pretrained(
@@ -33,9 +34,7 @@ stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__v
 stage_3.enable_model_cpu_offload()
 
 
-def predict(input_img):
-    prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'
-
+def predict(prompt):
     prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
     generator = torch.manual_seed(0)
     image = stage_1(
@@ -53,7 +52,7 @@ def predict(input_img):
     ).images
     image = stage_3(
         prompt=prompt, image=image, generator=generator, noise_level=100
-    ).images
+    ).images[0]
     return image
 
 
@@ -66,4 +65,4 @@ gradio_app = gr.Interface(
 )
 
 if __name__ == "__main__":
-    gradio_app.launch() # server_name="0.0.0.0"
+    gradio_app.launch(server_name="0.0.0.0") # server_name="0.0.0.0"

bapp.py

@@ -0,0 +1,78 @@
+import argparse
+from pathlib import Path
+
+import gradio as gr
+import torch
+from diffusers import DiffusionPipeline
+
+from visual_anagrams.views import get_views, VIEW_MAP_NAMES
+from visual_anagrams.samplers import sample_stage_1, sample_stage_2
+from visual_anagrams.utils import add_args, save_illusion, save_metadata
+
+stage_1 = DiffusionPipeline.from_pretrained(
+                "DeepFloyd/IF-I-M-v1.0",
+                variant="fp16",
+                torch_dtype=torch.float16)
+stage_2 = DiffusionPipeline.from_pretrained(
+                "DeepFloyd/IF-II-M-v1.0",
+                text_encoder=None,
+                variant="fp16",
+                torch_dtype=torch.float16,
+            )
+stage_1.enable_model_cpu_offload()
+stage_2.enable_model_cpu_offload()
+
+
+def generate_content(
+    style,
+    prompt_for_original,
+    prompt_for_transformed,
+    transformation,
+    num_inference_steps,
+    seed
+):
+    prompts = [prompt_for_original, prompt_for_transformed]
+    prompt_embeds = [stage_1.encode_prompt(f'{style} {p}'.strip()) for p in [prompts]]
+    prompt_embeds, negative_prompt_embeds = zip(*prompt_embeds)
+    prompt_embeds = torch.cat(prompt_embeds)
+    negative_prompt_embeds = torch.cat(negative_prompt_embeds)
+
+    views = ['identity', transformation]
+    views = get_views(views)
+
+    generator = torch.manual_seed(seed)
+    image = sample_stage_1(stage_1,
+                            prompt_embeds,
+                            negative_prompt_embeds,
+                            views,
+                            num_inference_steps=num_inference_steps,
+                            generator=generator)
+
+    image = sample_stage_2(stage_2,
+                           image,
+                           prompt_embeds,
+                           negative_prompt_embeds,
+                           views,
+                           num_inference_steps=num_inference_steps,
+                           generator=generator)
+    
+    return image, image_transformed, transformation_gif
+
+
+choices = list(VIEW_MAP_NAMES.keys())
+gradio_app = gr.Interface(
+    fn=generate_content,
+    inputs=[
+        gr.Textbox(label="Style", placeholder="an oil painting of"),
+        gr.Textbox(label="Prompt for original view", placeholder="a penguin"),
+        gr.Textbox(label="Prompt for transformed view", placeholder="a giraffe"),
+        gr.Dropdown(label="View transformation", choices=choices, value=choices[0]),
+        gr.Number(label="Number of diffusion steps", value=30, step=1, minimum=1, maximum=100),
+        gr.Number(label="Random seed", value=0, step=1, minimum=0, maximum=100000)
+    ],
+    outputs=[gr.Image(label="Illusion"), gr.Image(label="Original"), gr.Image(label="Transformed")],
+)
+
+
+if __name__ == "__main__":
+    gradio_app.launch(server_name="0.0.0.0") # server_name="0.0.0.0"

requirements.txt

@@ -1,8 +1,10 @@
 accelerate
 diffusers
+einops
 gradio
 safetensors
 sentencepiece
 transformers
 torch
+torchvision
 xformers

visual_anagrams/views/__init__.py

@@ -25,6 +25,19 @@ VIEW_MAP = {
     'inner_circle': InnerCircleView,
 }
 
+VIEW_MAP_NAMES = {
+    'Flip': 'flip',
+    'Rotate 90° clockwise': 'rotate_cw',
+    'Rotate 90° counter-clockwise': 'rotate_ccw',
+    'Rotate 180°': 'rotate_180',
+    'Invert colors': 'negate',
+    'Shear': 'skew',
+    'Patch permutation': 'patch_permute',
+    'Pixel permutation': 'pixel_permute',
+    'Jigsaw permutation': 'jigsaw',
+    'Rotate inner circle': 'inner_circle',
+}
+
 def get_views(view_names):
     '''
     Bespoke function to get views (just to make command line usage easier)