Update app.py

app.py

@@ -6,7 +6,6 @@ api.set_access_token(os.environ['HF_SECRET'])
 
 folder = HfFolder()
 folder.save_token(os.environ['HF_SECRET'])
-
 import math
 import time
 from threading import Lock
@@ -21,7 +20,7 @@ import torch
 from spacy import displacy
 
 from daam import trace
-from daam.utils import set_seed, cached_nlp
+from daam.utils import set_seed, cached_nlp, auto_autocast
 
 
 def dependency(text):
@@ -55,7 +54,8 @@ def get_args():
         'v1': 'runwayml/stable-diffusion-v1-5',
         'v2-base': 'stabilityai/stable-diffusion-2-base',
         'v2-large': 'stabilityai/stable-diffusion-2',
-        'v2-1-base': 'stabilityai/stable-diffusion-2-1-base'
+        'v2-1-base': 'stabilityai/stable-diffusion-2-1-base',
+        'v2-1-large': 'stabilityai/stable-diffusion-2-1',
     }
 
     parser = argparse.ArgumentParser()
@@ -93,7 +93,7 @@ def main():
             new_prompt = ' '.join(new_prompt)
 
         merge_idxs, words = get_tokenizing_mapping(prompt, pipe.tokenizer)
-        with torch.cuda.amp.autocast(dtype=torch.float16), lock:
+        with auto_autocast(dtype=torch.float16), lock:
             try:
                 plt.close('all')
                 plt.clf()
@@ -167,10 +167,8 @@ def main():
 
     with gr.Blocks(css='scrollbar.css') as demo:
         md = '''# DAAM: Attention Maps for Interpreting Stable Diffusion
-        Check out the **new** paper (2022/12/7): [What the DAAM: Interpreting Stable Diffusion Using Cross Attention](http://arxiv.org/abs/2210.04885).
+        Check out the paper: [What the DAAM: Interpreting Stable Diffusion Using Cross Attention](http://arxiv.org/abs/2210.04885).
         See our (much cleaner) [DAAM codebase](https://github.com/castorini/daam) on GitHub.
-
-        Current model: Stable Diffusion 2.1 base.
         '''
         gr.Markdown(md)
 
@@ -220,7 +218,7 @@ def main():
 
     while True:
         try:
-            demo.launch()
+            demo.launch(server_name='0.0.0.0', server_port=args.port)
         except OSError:
             gr.close_all()
         except KeyboardInterrupt: