finish

parti_prompts.py

@@ -35,15 +35,16 @@ def get_karlo_eval(ckpt):
     return karlo_eval
 
 def get_if_eval(ckpt):
-    pipe_low = DiffusionPipeline.from_pretrained(ckpt, safety_checker=None, torch_dtype=torch.float16)
+    pipe_low = DiffusionPipeline.from_pretrained(ckpt, safety_checker=None, watermarker=None, torch_dtype=torch.float16, variant="fp16")
     pipe_low.enable_model_cpu_offload()
 
-    pipe_up = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", safety_checker=None, text_encoder=pipe_low.text_encoder, torch_dtype=torch.float16)
+    pipe_up = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", safety_checker=None, watermarker=None, text_encoder=pipe_low.text_encoder, torch_dtype=torch.float16, variant="fp16")
     pipe_up.enable_model_cpu_offload()
 
     def if_eval(prompt, generator=None):
-        images = pipe_low(prompt, num_inference_steps=NUM_INFERENCE_STEPS, generator=generator, output_type="pt").images
-        images = pipe_up(promtp=prompt, images=images, num_inference_steps=NUM_INFERENCE_STEPS).images
+        prompt_embeds, negative_prompt_embeds = pipe_low.encode_prompt(prompt)
+        images = pipe_low(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=NUM_INFERENCE_STEPS, generator=generator, output_type="pt").images
+        images = pipe_up(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, image=images, num_inference_steps=NUM_INFERENCE_STEPS, generator=generator).images
         return images
 
     return if_eval
@@ -69,6 +70,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     dataset = load_dataset("nateraw/parti-prompts")["train"]
+    # dataset = dataset.select(range(4))
 
     eval_fn = MODELS[args.model_repo_or_id](args.model_repo_or_id)