Update app.py

app.py

@@ -20,8 +20,7 @@ safety_model_id = "CompVis/stable-diffusion-safety-checker"
 safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
 safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
 
-# Change device to CPU
-device = "cpu"
+device = "cuda"
 
 noise_scheduler = DDIMScheduler(
     num_train_timesteps=1000,
@@ -32,10 +31,10 @@ noise_scheduler = DDIMScheduler(
     set_alpha_to_one=False,
     steps_offset=1,
 )
-vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float32)  # Change dtype to float32 for CPU compatibility
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
 pipe = StableDiffusionPipeline.from_pretrained(
     base_model_path,
-    torch_dtype=torch.float32,  # Change dtype to float32 for CPU
+    torch_dtype=torch.float16,
     scheduler=noise_scheduler,
     vae=vae,
     feature_extractor=safety_feature_extractor,
@@ -53,7 +52,7 @@ app.prepare(ctx_id=0, det_size=(640, 640))
 
 cv2.setNumThreads(1)
 
-@spaces.GPU(enable_queue=True)  # Changed to GPU to allow flexibility between CPU/GPU
+@spaces.GPU(enable_queue=True)
 def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, progress=gr.Progress(track_tqdm=True)):
     faceid_all_embeds = []
     first_iteration = True
@@ -63,7 +62,7 @@ def generate_image(images, prompt, negative_prompt, preserve_face_structure, fac
         faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
         faceid_all_embeds.append(faceid_embed)
         if(first_iteration and preserve_face_structure):
-            face_image = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224)  # you can also segment the face
+            face_image = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224) # you can also segment the face
             first_iteration = False
             
     average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
@@ -96,11 +95,9 @@ def swap_to_gallery(images):
 
 def remove_back_to_files():
     return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
-
 css = '''
 h1{margin-bottom: 0 !important}
 '''
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# IP-Adapter-FaceID Plus demo")
     gr.Markdown("Demo for the [h94/IP-Adapter-FaceID model](https://huggingface.co/h94/IP-Adapter-FaceID) - Generate AI images with your own face - Non-commercial license")
@@ -137,4 +134,4 @@ with gr.Blocks(css=css) as demo:
     
     gr.Markdown("This demo includes extra features to mitigate the implicit bias of the model and prevent explicit usage of it to generate content with faces of people, including third parties, that is not safe for all audiences, including naked or semi-naked people.")
     
-demo.launch()
+demo.launch()