Upload control.py

control.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+from fal_serverless import isolated, cached
+
+from pathlib import Path
+import base64
+import io
+
+requirements = [
+    "controlnet-aux",
+    "diffusers",
+    "torch",
+    "mediapipe",
+    "transformers",
+    "accelerate",
+    "xformers"
+]
+
+
+def get_image_from_url_as_bytes(url: str) -> bytes:
+    import requests
+
+    response = requests.get(url)
+    # This will raise an exception if the request returned an HTTP error code
+    response.raise_for_status()
+    return response.content
+
+def read_image_bytes(file_path):
+    with open(file_path, "rb") as file:
+        image_bytes = file.read()
+    return image_bytes
+
+@cached
+def load_model():
+    import torch
+    from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+
+    controlnet = ControlNetModel.from_pretrained(
+        "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16
+    )
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        "peterwilli/deliberate-2", controlnet=controlnet, torch_dtype=torch.float16
+    )
+
+    pipe = pipe.to("cuda:0")
+    pipe.unet.to(memory_format=torch.channels_last)
+    pipe.controlnet.to(memory_format=torch.channels_last)
+    return pipe
+
+
+def resize_image(input_image, resolution):
+    import cv2
+    import numpy as np
+
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(
+        input_image,
+        (W, H),
+        interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA,
+    )
+    return img
+
+@isolated(
+    requirements=requirements,
+    machine_type="GPU",
+    keep_alive=30,
+    serve=True
+)
+def generate(
+    image_url: str, prompt: str, num_samples: int, num_steps: int, gcs=False
+) -> list[bytes] | None:
+
+    from controlnet_aux import CannyDetector
+    from PIL import Image
+    import numpy as np
+    import uuid
+    import os
+    from base64 import b64encode
+
+    image_bytes = get_image_from_url_as_bytes(image_url)
+
+    pipe = load_model()
+    image = Image.open(io.BytesIO(image_bytes))
+
+    canny = CannyDetector()
+    init_image = image.convert("RGB")
+
+    init_image = resize_image(np.asarray(init_image), 512)
+    detected_map = canny(init_image, 100, 200)
+    image = Image.fromarray(detected_map)
+
+    negative_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+    results = pipe(
+        prompt=prompt,
+        image=image,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_steps,
+        num_images_per_prompt=num_samples
+    ).images
+
+    result_id = uuid.uuid4()
+    out_dir = Path(f"/data/cn-results/{result_id}")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+
+    for i, res in enumerate(results):
+        res.save(out_dir / f"res_{i}.png")
+
+    file_names = [
+        f for f in os.listdir(out_dir) if os.path.isfile(os.path.join(out_dir, f))
+    ]
+
+    list_of_bytes = [read_image_bytes(out_dir / f) for f in file_names]
+    raw_image = list_of_bytes[0]
+
+    return b64encode(raw_image).decode("utf-8")