envs

app.py

@@ -81,9 +81,12 @@ If you have any questions, please feel free to reach me out at <b>ywl@stu.pku.ed
 # """
 
 os.makedirs("models/personalized")
+os.makedirs("models/sd1-5")
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/flow_controlnet.ckpt -P models/')
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/image_controlnet.ckpt -P models/')
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/unet.ckpt -P models/')
+os.system(f'wget https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/config.json -P models/sd1-5/')
+os.system(f'wget https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/diffusion_pytorch_model.bin -P models/sd1-5/')
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/helloobjects_V12c.safetensors -P models/personalized')
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/TUSUN.safetensors -P models/personalized')
 
@@ -194,7 +197,7 @@ class ImageConductor:
         text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
         vae          = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
         inference_config = OmegaConf.load("configs/inference/inference.yaml")
-        unet = UNet3DConditionFlowModel.from_pretrained_2d("runwayml/stable-diffusion-v1-5", subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
+        unet = UNet3DConditionFlowModel.from_pretrained_2d("models/sd1-5/", subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
 
         self.vae = vae