remove custom mixins

huggingface · Sep 13, 2024 · a706204 · a706204
1 parent 53eedc6
commit a706204
Show file tree

Hide file tree

Showing 10 changed files with 233 additions and 4,538 deletions.
diff --git a/optimum/onnxruntime/modeling_diffusion.py b/optimum/onnxruntime/modeling_diffusion.py
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
@@ -401,3 +401,23 @@ def evaluation_loop(
         metrics = {}
 
     return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=len(dataset))
+
+
+def np_to_pt(np_object, device):
+    if isinstance(np_object, np.ndarray):
+        if np_object.ndim == 4:
+            return torch.from_numpy(np_object).permute(0, 3, 1, 2)
+        elif np_object.ndim == 3:
+            return torch.from_numpy(np_object).permute(2, 0, 1)
+        else:
+            return torch.from_numpy(np_object)
+    elif isinstance(np_object, list) and isinstance(np_object[0], np.ndarray):
+        return [np_to_pt(a, device) for a in np_object]
+    elif isinstance(np_object, dict) and isinstance(next(iter(np_object.values())), np.ndarray):
+        return {k: np_to_pt(v, device) for k, v in np_object.items()}
+    elif isinstance(np_object, np.random.RandomState):
+        return torch.Generator(device=device).manual_seed(int(np_object.get_state()[1][0]))
+    elif isinstance(np_object, list) and isinstance(np_object[0], np.random.RandomState):
+        return [torch.Generator(device=device).manual_seed(int(a.get_state()[1][0])) for a in np_object]
+    else:
+        return np_object
diff --git a/optimum/pipelines/diffusers/pipeline_latent_consistency.py b/optimum/pipelines/diffusers/pipeline_latent_consistency.py
diff --git a/optimum/pipelines/diffusers/pipeline_stable_diffusion.py b/optimum/pipelines/diffusers/pipeline_stable_diffusion.py
diff --git a/optimum/pipelines/diffusers/pipeline_stable_diffusion_img2img.py b/optimum/pipelines/diffusers/pipeline_stable_diffusion_img2img.py
diff --git a/optimum/pipelines/diffusers/pipeline_stable_diffusion_inpaint.py b/optimum/pipelines/diffusers/pipeline_stable_diffusion_inpaint.py
diff --git a/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py b/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py
diff --git a/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl_img2img.py b/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl_img2img.py
diff --git a/optimum/pipelines/diffusers/pipeline_utils.py b/optimum/pipelines/diffusers/pipeline_utils.py
diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py
@@ -240,10 +240,20 @@ def test_negative_prompt(self, model_arch: str):
                 inputs["negative_prompt_embeds"],
                 inputs["pooled_prompt_embeds"],
                 inputs["negative_pooled_prompt_embeds"],
-            ) = pipeline.encode_prompt(prompt=prompt, negative_prompt=negative_prompt)
+            ) = pipeline.encode_prompt(
+                prompt=prompt,
+                num_images_per_prompt=1,
+                device=torch.device("cpu"),
+                do_classifier_free_guidance=True,
+                negative_prompt=negative_prompt,
+            )
         else:
             inputs["prompt_embeds"], inputs["negative_prompt_embeds"] = pipeline.encode_prompt(
-                prompt=prompt, negative_prompt=negative_prompt
+                prompt=prompt,
+                num_images_per_prompt=1,
+                device=torch.device("cpu"),
+                do_classifier_free_guidance=True,
+                negative_prompt=negative_prompt,
             )
 
         images_2 = pipeline(**inputs, generator=get_generator("pt", SEED)).images