Merge branch 'main' of https://huggingface.co./diffusers/tools

clear_mem.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+import torch
+import gc
+
+shape = (10,000)
+
+input = torch.ones((shape, shape), device="cuda")
+
+
+def clear_memory(model):
+    model.to('cpu')
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+    torch.clear_autocast_cache()
+
+for _ in range(6):
+    linear = torch.nn.Linear(shape, shape).to("cuda")
+    output = linear(input)
+
+    clear_memory(linear)

run_sdxl_dummy.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+import torch
+import numpy as np
+from diffusers import StableDiffusionXLPipeline
+
+path = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+
+pipe = StableDiffusionXLPipeline.from_pretrained(path)
+pipe.unet.set_default_attn_processor()
+
+prompt = "An astronaut riding a green horse on Mars"
+steps = 3
+
+batch_size, height, width, ch = 1, 32, 32, 4
+num_elems = batch_size * height * width * ch
+latents = (torch.arange(num_elems) / num_elems)[:, None, None, None].reshape(batch_size, ch, width, height)
+print("latents", latents.abs().sum())
+
+image = pipe(prompt, latents=latents, num_inference_steps=3, output_type="np", guidance_scale=7.5).images[0]
+
+print(np.abs(image).sum())

run_sdxl_dummy_flax.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+from diffusers import FlaxStableDiffusionXLPipeline
+import numpy as np
+import jax.numpy as jnp
+import jax
+
+path = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+
+pipe, params = FlaxStableDiffusionXLPipeline.from_pretrained(path)
+
+prompt = "An astronaut riding a green horse on Mars"
+steps = 3
+
+batch_size, height, width, ch = 1, 32, 32, 4
+num_elems = batch_size * height * width * ch
+rng = jax.random.PRNGKey(0)
+latents = (jnp.arange(num_elems) / num_elems)[:, None, None, None].reshape(batch_size, ch, width, height)
+
+print("latents", np.abs(np.asarray(latents)).sum())
+
+prompt_embeds = pipe.prepare_inputs(prompt)
+
+image = pipe(prompt_embeds, params, rng, latents=latents, num_inference_steps=3, output_type="np").images[0]
+
+print(np.abs(np.asarray(image)).sum())