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)