add token

app.py

@@ -2,6 +2,7 @@ from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeli
 import gradio as gr
 import torch
 from PIL import Image
+import os
 
 scheduler = DPMSolverMultistepScheduler(
     beta_start=0.00085,
@@ -56,16 +57,18 @@ last_mode = "txt2img"
 current_model = models[1]
 current_model_path = current_model.path
 
+auth_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+
 if is_colab:
-  pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler)
+  pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
 
 else: # download all models
-  vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16)
+  vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16, use_auth_token=auth_token)
   for model in models[1:]:
     try:
-        unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
-        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
-        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
+        unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16, use_auth_token=auth_token)
+        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
+        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
     except:
         models.remove(model)
   pipe = models[1].pipe_t2i
@@ -104,7 +107,7 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
         current_model_path = model_path
 
         if is_colab or current_model == models[0]:
-          pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
+          pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, use_auth_token=auth_token)
         else:
           pipe.to("cpu")
           pipe = current_model.pipe_t2i
@@ -135,7 +138,7 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
         current_model_path = model_path
 
         if is_colab or current_model == models[0]:
-          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
+          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, use_auth_token=auth_token)
         else:
           pipe.to("cpu")
           pipe = current_model.pipe_i2i
@@ -266,19 +269,21 @@ with gr.Blocks(css=css) as demo:
     prompt.submit(inference, inputs=inputs, outputs=image_out)
     generate.click(inference, inputs=inputs, outputs=image_out)
 
+
+# TODO: the docs here are wrong.
     ex = gr.Examples([
-        [models[1].name, "jason bateman disassembling the demon core", 7.5, 50],
-        [models[4].name, "portrait of dwayne johnson", 7.0, 75],
-        [models[5].name, "portrait of a beautiful alyx vance half life", 10, 50],
-        [models[6].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
-        [models[5].name, "fantasy portrait painting, digital art", 4.0, 30],
+        [models[1+2].name, "jason bateman disassembling the demon core", 7.5, 50],
+        [models[4+2].name, "portrait of dwayne johnson", 7.0, 75],
+        [models[5+2].name, "portrait of a beautiful alyx vance half life", 10, 50],
+        [models[6+2].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
+        [models[5+2].name, "fantasy portrait painting, digital art", 4.0, 30],
     ], [model_name, prompt, guidance, steps, seed], image_out, inference, cache_examples=False)
 
     gr.Markdown('''
       Models by [@nitrosocke](https://huggingface.co/nitrosocke), [@haruu1367](https://twitter.com/haruu1367), [@Helixngc7293](https://twitter.com/DGSpitzer) and others. ❤️<br>
       Space by: [![Twitter Follow](https://img.shields.io/twitter/follow/hahahahohohe?label=%40anzorq&style=social)](https://twitter.com/hahahahohohe)
         
-      ![visitors](https://visitor-badge.glitch.me/badge?page_id=anzorq.finetuned_diffusion)
+      ![visitors](https://visitor-badge.glitch.me/badge?page_id=LuChengTHU.dpmsolver_sdm)
     ''')
 
 if not is_colab: