envs

app.py

@@ -86,6 +86,7 @@ os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/imag
 os.system(f'wget https://huggingface.co/TencentARC/ImageConductor/blob/main/unet.ckpt -P models/')
 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')
+os.system(f'git clone https://huggingface.co/runwayml/stable-diffusion-v1-5')
 
 
 
@@ -190,11 +191,11 @@ def points_to_flows(track_points, model_length, height, width):
 class ImageConductor:
     def __init__(self, device, unet_path, image_controlnet_path, flow_controlnet_path, height, width, model_length, lora_rank=64):
         self.device = device
-        tokenizer    = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
-        text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
-        vae          = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
+        tokenizer    = CLIPTokenizer.from_pretrained("stable-diffusion-v1-5", subfolder="tokenizer")
+        text_encoder = CLIPTextModel.from_pretrained("stable-diffusion-v1-5", subfolder="text_encoder").cuda()
+        vae          = AutoencoderKL.from_pretrained("stable-diffusion-v1-5", subfolder="vae").cuda()
         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("stable-diffusion-v1-5", subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
 
         self.vae = vae