init

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+import sys
+# sys.path.append("LaVi-Bridge/test")
+# from llama2_unet_diffusion_lens import call_diffusion_lens
+from diffusion_lens import get_images
+
+import os
+import subprocess
+
+def prepare_images(images):
+    return images
+
+
+# def call_diffusion_lens(prompt):
+#     # os.chdir('LaVi-Bridge/test')
+#     command = f"python -u llama2_unet.py --ckpt_dir 'LaVi-Bridge/llama2_unet' --output_dir 'output' --llama2_dir 'meta-llama/Llama-2-7b-hf' --prompt '{prompt}'"
+#     subprocess.run(command, shell=True)
+#     return 'output'
+
+
+def get_prompt(prompt):
+    print('prompt:', prompt)
+    print('calling diffusion lens')
+
+    # parser.add_argument("--ckpt_dir", type=str, default="")
+    # parser.add_argument("--output_dir", type=str, default="")
+    # parser.add_argument("--llama2_dir", type=str, default="")
+    # parser.add_argument("--prompts_path", type=str, default="inputs/in.txt")
+    # parser.add_argument("--use_chat", action="store_true")
+    # parser.add_argument("--generate_text", action="store_true")
+    # parser.add_argument("--dont_use_lora", action="store_true")
+    # parser.add_argument("--is_gradio", action="store_true")
+
+    # args = {
+    #     'ckpt_dir': 'LaVi-Bridge/LaVi-Bridge/llama2_unet',
+    #     'output_dir': 'output',
+    #     'llama2_dir': 'meta-llama/Llama-2-7b-hf',
+    #     'prompt': prompt,
+    #     'use_chat': False,
+    #     'generate_text': False,
+    #     'dont_use_lora': False,
+    #     'is_gradio': True,
+    #     'prompts_path': None
+    # }
+    # images = call_diffusion_lens(args, prompt)
+    # print('done calling diffusion lens')
+    # print('number of images:', len(images))
+    # images = prepare_images(images)
+    # print('done preparing images')
+    image = get_images(prompt)
+    return image
+
+if __name__ == '__main__':
+    print('starting')
+    get_prompt("A photo of a cat")
+    print('done')
+
+
+# iface = gr.Interface(fn=get_prompt, inputs="text", outputs="image", title="Diffusion Lens")
+# iface.launch()
+
+import gradio as gr
+
+def display_images(images):
+    # Prepare images for display
+    return [gr.Image(image) for image in images]
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        gallery = gr.Gallery(
+            label="Generated images", show_label=False, elem_id="gallery",
+            columns=[1], rows=[1], object_fit="contain", height="auto")
+        btn = gr.Button("Generate images", scale=0)
+
+        btn.click(get_prompt, 'text', gallery)
+
+    demo.launch()

diffusion_lens.py

@@ -0,0 +1,8 @@
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
+
+pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
+
+def get_images(prompt, skip_layers=0):
+    images = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1)
+    return images

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+diffusers
+gradio