refactor and bugfix

app.py

@@ -8,38 +8,7 @@ from insightface.utils import face_align
 import gradio as gr
 from huggingface_hub import hf_hub_download
 
-base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
-vae_model_path = "stabilityai/sd-vae-ft-mse"
-image_encoder_path = "IP-Adapter/models/image_encoder"
-ip_ckpt = "IP-Adapter-FaceID/ip-adapter-faceid-plus_sd15.bin"
-
-
-if torch.cuda.is_available():
-    device = 'cuda'
-    torch_dtype = torch.float16
-else:
-    device = 'cpu'
-    torch_dtype = torch.float32
-print(f'Using device: {device}')
-
-noise_scheduler = DDIMScheduler(
-    num_train_timesteps=1000,
-    beta_start=0.00085,
-    beta_end=0.012,
-    beta_schedule="scaled_linear",
-    clip_sample=False,
-    set_alpha_to_one=False,
-    steps_offset=1,
-)
-vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch_dtype)
-pipe = StableDiffusionPipeline.from_pretrained(
-    base_model_path,
-    torch_dtype=torch_dtype,
-    scheduler=noise_scheduler,
-    vae=vae,
-    feature_extractor=None,
-    safety_checker=None
-)
+
 
 def download_models():
     hf_hub_download(
@@ -56,10 +25,45 @@ def download_models():
         local_dir='IP-Adapter')
 
 
-ip_model = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_ckpt, device, num_tokens=4, torch_dtype=torch_dtype)
+def get_ip_model():
+    download_models()
+    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
+    vae_model_path = "stabilityai/sd-vae-ft-mse"
+    image_encoder_path = "IP-Adapter/models/image_encoder"
+    ip_ckpt = "IP-Adapter-FaceID/ip-adapter-faceid-plus_sd15.bin"
+
+    if torch.cuda.is_available():
+        device = 'cuda'
+        torch_dtype = torch.float16
+    else:
+        device = 'cpu'
+        torch_dtype = torch.float32
+    print(f'Using device: {device}')
+
+    noise_scheduler = DDIMScheduler(
+        num_train_timesteps=1000,
+        beta_start=0.00085,
+        beta_end=0.012,
+        beta_schedule="scaled_linear",
+        clip_sample=False,
+        set_alpha_to_one=False,
+        steps_offset=1,
+    )
+    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch_dtype)
+    pipe = StableDiffusionPipeline.from_pretrained(
+        base_model_path,
+        torch_dtype=torch_dtype,
+        scheduler=noise_scheduler,
+        vae=vae,
+        feature_extractor=None,
+        safety_checker=None
+    )
 
+    ip_model = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_ckpt, device, num_tokens=4, torch_dtype=torch_dtype)
+    return ip_model
 
 
+ip_model = get_ip_model()
 app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.2)