def send_requests()

in ai-ml/gke-ray/rayserve/stable-diffusion/stable_diffusion_tpu_req.py [0:0]

• 26 lines of code
• 4 McCabe index (conditional complexity)

def send_requests(num_requests: int, batch_size: int, save_pictures: bool,
                  url: str = "http://localhost:8000/imagine"):
  """Sends a list of requests and processes the responses."""
  print("num_requests: ", num_requests)
  print("batch_size: ", batch_size)
  print("url: ", url)
  print("save_pictures: ", save_pictures)

  prompts = _PROMPTS
  if num_requests > len(_PROMPTS):
    # Repeat until larger than num_requests
    prompts = _PROMPTS * int(np.ceil(num_requests / len(_PROMPTS)))

  prompts = np.random.choice(
      prompts, num_requests, replace=False)

  with futures.ThreadPoolExecutor(max_workers=batch_size) as executor:
    raw_images = list(
        tqdm(
            executor.map(
                functools.partial(send_request_and_receive_image, url=url),
                prompts,
            ),
            total=len(prompts),
        )
    )

  if save_pictures:
    print("Saving pictures to diffusion_results.png")
    images = [Image.open(raw_image) for raw_image in raw_images]
    grid = image_grid(images, 2, num_requests // 2)
    grid.save("./diffusion_results.png")