Update app filename

app.py

@@ -1,68 +1,92 @@
-from diffusers import DiffusionPipeline
-from diffusers.utils import pt_to_pil
+import argparse
+from pathlib import Path
+
 import gradio as gr
 import torch
-import numpy as np
+from diffusers import DiffusionPipeline
+from icecream import ic
 
+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
+from visual_anagrams.animate import animate_two_view
 
 stage_1 = DiffusionPipeline.from_pretrained(
-    "DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float16
-)
-stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
-stage_1.enable_model_cpu_offload()
+                "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_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+                "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()
 
-# stage 3
-safety_modules = {
-    "feature_extractor": stage_1.feature_extractor,
-    "safety_checker": stage_1.safety_checker,
-    "watermarker": stage_1.watermarker,
-}
-stage_3 = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-x4-upscaler",
-    **safety_modules,
-    torch_dtype=torch.float16
-)
-stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
-stage_3.enable_model_cpu_offload()
 
+def generate_content(
+    style,
+    prompt_for_original,
+    prompt_for_transformed,
+    transformation,
+    num_inference_steps,
+    seed
+):
+    prompts = [f'{style} {p}'.strip() for p in [prompt_for_original, prompt_for_transformed]]
+    prompt_embeds = [stage_1.encode_prompt(p) 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', VIEW_MAP_NAMES[transformation]]
+    views = get_views(views)
 
-def predict(prompt):
-    prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
-    generator = torch.manual_seed(0)
-    image = stage_1(
-        prompt_embeds=prompt_embeds,
-        negative_prompt_embeds=negative_embeds,
-        generator=generator,
-        output_type="pt",
-    ).images
-    image = stage_2(
-        image=image,
-        prompt_embeds=prompt_embeds,
-        negative_prompt_embeds=negative_embeds,
-        generator=generator,
-        output_type="pt",
-    ).images
-    image = stage_3(
-        prompt=prompt, image=image, generator=generator, noise_level=100
-    ).images[0]
-    return image
+    generator = torch.manual_seed(seed)
 
+    print("Sample stage 1")
+    image = sample_stage_1(stage_1,
+                           prompt_embeds,
+                           negative_prompt_embeds,
+                           views,
+                           num_inference_steps=num_inference_steps,
+                           generator=generator)
 
+    print("Sample stage 2")
+    image = sample_stage_2(stage_2,
+                           image,
+                           prompt_embeds,
+                           negative_prompt_embeds,
+                           views,
+                           num_inference_steps=num_inference_steps,
+                           generator=generator)
+    save_illusion(image, views, Path(""))
+    
+    size = image.shape[-1]
+    animate_two_view(
+        f"sample_{size}.png",
+        views[1],
+        prompts[0],
+        prompts[1],
+    )
+    return 'tmp.mp4', f"sample_{size}.png", f"sample_{size}.views.png"
+
+
+choices = list(VIEW_MAP_NAMES.keys())
 gradio_app = gr.Interface(
-    fn=predict,
-    inputs="text",
-    outputs="image",
-    title="Text to Image Generator",
-    description="Enter a text string to generate an image.",
+    fn=generate_content,
+    title="Multi-View Illusion Diffusion",
+    inputs=[
+        gr.Textbox(label="Style", placeholder="an oil painting of"),
+        gr.Textbox(label="Prompt for original view", placeholder="a dress"),
+        gr.Textbox(label="Prompt for transformed view", placeholder="an old man"),
+        gr.Dropdown(label="View transformation", choices=choices, value=choices[0]),
+        gr.Number(label="Number of diffusion steps", value=75, step=1, minimum=1, maximum=300),
+        gr.Number(label="Random seed", value=0, step=1, minimum=0, maximum=100000)
+    ],
+    outputs=[gr.Video(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"
+    gradio_app.launch() # server_name="0.0.0.0"

bapp.py

@@ -1,92 +0,0 @@
-import argparse
-from pathlib import Path
-
-import gradio as gr
-import torch
-from diffusers import DiffusionPipeline
-from icecream import ic
-
-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
-from visual_anagrams.animate import animate_two_view
-
-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 = [f'{style} {p}'.strip() for p in [prompt_for_original, prompt_for_transformed]]
-    prompt_embeds = [stage_1.encode_prompt(p) 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', VIEW_MAP_NAMES[transformation]]
-    views = get_views(views)
-
-    generator = torch.manual_seed(seed)
-
-    print("Sample stage 1")
-    image = sample_stage_1(stage_1,
-                           prompt_embeds,
-                           negative_prompt_embeds,
-                           views,
-                           num_inference_steps=num_inference_steps,
-                           generator=generator)
-
-    print("Sample stage 2")
-    image = sample_stage_2(stage_2,
-                           image,
-                           prompt_embeds,
-                           negative_prompt_embeds,
-                           views,
-                           num_inference_steps=num_inference_steps,
-                           generator=generator)
-    save_illusion(image, views, Path(""))
-    
-    size = image.shape[-1]
-    animate_two_view(
-        f"sample_{size}.png",
-        views[1],
-        prompts[0],
-        prompts[1],
-    )
-    return 'tmp.mp4', f"sample_{size}.png", f"sample_{size}.views.png"
-
-
-choices = list(VIEW_MAP_NAMES.keys())
-gradio_app = gr.Interface(
-    fn=generate_content,
-    title="Multi-View Illusion Diffusion",
-    inputs=[
-        gr.Textbox(label="Style", placeholder="an oil painting of"),
-        gr.Textbox(label="Prompt for original view", placeholder="a dress"),
-        gr.Textbox(label="Prompt for transformed view", placeholder="an old man"),
-        gr.Dropdown(label="View transformation", choices=choices, value=choices[0]),
-        gr.Number(label="Number of diffusion steps", value=50, step=1, minimum=1, maximum=300),
-        gr.Number(label="Random seed", value=0, step=1, minimum=0, maximum=100000)
-    ],
-    outputs=[gr.Video(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"

test_app.py

@@ -0,0 +1,68 @@
+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(
+    "DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float16
+)
+stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_1.enable_model_cpu_offload()
+stage_2 = DiffusionPipeline.from_pretrained(
+    "DeepFloyd/IF-II-M-v1.0",
+    text_encoder=None,
+    variant="fp16",
+    torch_dtype=torch.float16,
+)
+stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_2.enable_model_cpu_offload()
+
+# stage 3
+safety_modules = {
+    "feature_extractor": stage_1.feature_extractor,
+    "safety_checker": stage_1.safety_checker,
+    "watermarker": stage_1.watermarker,
+}
+stage_3 = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-x4-upscaler",
+    **safety_modules,
+    torch_dtype=torch.float16
+)
+stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_3.enable_model_cpu_offload()
+
+
+def predict(prompt):
+    prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
+    generator = torch.manual_seed(0)
+    image = stage_1(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_embeds,
+        generator=generator,
+        output_type="pt",
+    ).images
+    image = stage_2(
+        image=image,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_embeds,
+        generator=generator,
+        output_type="pt",
+    ).images
+    image = stage_3(
+        prompt=prompt, image=image, generator=generator, noise_level=100
+    ).images[0]
+    return image
+
+
+gradio_app = gr.Interface(
+    fn=predict,
+    inputs="text",
+    outputs="image",
+    title="Text to Image Generator",
+    description="Enter a text string to generate an image.",
+)
+
+if __name__ == "__main__":
+    gradio_app.launch(server_name="0.0.0.0") # server_name="0.0.0.0"