Update app.py

app.py

@@ -1,3 +1,34 @@
+import torch
+import types
+torch.cuda.get_device_capability = lambda *args, **kwargs: (8, 6)
+torch.cuda.get_device_properties = lambda *args, **kwargs: types.SimpleNamespace(name='NVIDIA A10G', major=8, minor=6, total_memory=23836033024, multi_processor_count=80)
+
+import huggingface_hub
+huggingface_hub.snapshot_download(
+    repo_id='camenduru/PASD',
+    allow_patterns=[
+        'pasd/**',
+        'pasd_light/**',
+        'pasd_light_rrdb/**',
+        'pasd_rrdb/**',
+    ],
+    local_dir='PASD/runs',
+    local_dir_use_symlinks=False,
+)
+huggingface_hub.hf_hub_download(
+    repo_id='camenduru/PASD',
+    filename='majicmixRealistic_v6.safetensors',
+    local_dir='PASD/checkpoints/personalized_models',
+    local_dir_use_symlinks=False,
+)
+huggingface_hub.hf_hub_download(
+    repo_id='akhaliq/RetinaFace-R50',
+    filename='RetinaFace-R50.pth',
+    local_dir='PASD/annotator/ckpts',
+    local_dir_use_symlinks=False,
+)
+
+import sys; sys.path.append('./PASD')
 import spaces
 import os
 import datetime
@@ -32,10 +63,10 @@ else:
     from models.pasd.unet_2d_condition import UNet2DConditionModel
     from models.pasd.controlnet import ControlNetModel
 
-pretrained_model_path = "checkpoints/stable-diffusion-v1-5"
-ckpt_path = "runs/pasd/checkpoint-100000"
+pretrained_model_path = "runwayml/stable-diffusion-v1-5"
+ckpt_path = "PASD/runs/pasd/checkpoint-100000"
 #dreambooth_lora_path = "checkpoints/personalized_models/toonyou_beta3.safetensors"
-dreambooth_lora_path = "checkpoints/personalized_models/majicmixRealistic_v6.safetensors"
+dreambooth_lora_path = "PASD/checkpoints/personalized_models/majicmixRealistic_v6.safetensors"
 #dreambooth_lora_path = "checkpoints/personalized_models/Realistic_Vision_V5.1.safetensors"
 weight_dtype = torch.float16
 device = "cuda"
@@ -44,7 +75,7 @@ scheduler = UniPCMultistepScheduler.from_pretrained(pretrained_model_path, subfo
 text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
 tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
 vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
-feature_extractor = CLIPImageProcessor.from_pretrained(f"{pretrained_model_path}/feature_extractor")
+feature_extractor = CLIPImageProcessor.from_pretrained(pretrained_model_path, subfolder="feature_extractor")
 unet = UNet2DConditionModel.from_pretrained(ckpt_path, subfolder="unet")
 controlnet = ControlNetModel.from_pretrained(ckpt_path, subfolder="controlnet")
 vae.requires_grad_(False)
@@ -191,7 +222,7 @@ with gr.Blocks(css=css) as demo:
         """)
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(type="filepath", sources=["upload"], value="samples/frog.png")
+                input_image = gr.Image(type="filepath", sources=["upload"], value="PASD/samples/frog.png")
                 prompt_in = gr.Textbox(label="Prompt", value="Frog")
                 with gr.Accordion(label="Advanced settings", open=False):
                     added_prompt = gr.Textbox(label="Added Prompt", value='clean, high-resolution, 8k, best quality, masterpiece')