Added ControlNet ComfyUI workflow

blender/.beamignore

@@ -0,0 +1,28 @@
+# Generated by Beam SDK
+.beamignore
+pyproject.toml
+.git
+.idea
+.python-version
+.vscode
+.venv
+venv
+__pycache__
+.DS_Store
+.config
+drive/MyDrive
+.coverage
+.pytest_cache
+.ipynb
+.ruff_cache
+.dockerignore
+.ipynb_checkpoints
+.env.local
+.envrc
+**/__pycache__/
+**/.pytest_cache/
+**/node_modules/
+**/.venv/
+*.pyc
+.next/
+.circleci

blender/app.py

@@ -0,0 +1,63 @@
+"""
+This script uses Blender to render a 3D scene. It downloads a .blend file, executes a Python script to render the scene,
+
+To run this on Beam, you can run 'python app.py' in the terminal. Change the render_test.json and render.py files to your own use case.
+"""
+
+from beam import function, Image, Output
+from pathlib import Path
+import subprocess
+
+blender_image = (
+    Image(python_version="python3.11")
+    .add_python_packages(["bpy"])
+    .add_commands(
+        [
+            "apt update && apt install -y xorg libxkbcommon0",
+            "wget -q https://download.blender.org/release/Blender4.0/blender-4.0.2-linux-x64.tar.xz -O /tmp/blender.tar.xz",
+            "tar -xf /tmp/blender.tar.xz -C /opt",
+            "mv /opt/blender-* /opt/blender",
+            "ln -s /opt/blender/blender /usr/local/bin/blender",
+        ]
+    )
+)
+
+
+@function(image=blender_image, cpu=12, memory="32Gi", gpu="A10G")
+def render(script_content, json_content, output_name):
+    blend_path = "/tmp/render_test.blend"
+    script_path = "/tmp/render_alternative.py"
+    json_path = "/tmp/render_test.json"
+    output_path = f"/tmp/{output_name}"
+
+    subprocess.run(
+        "wget https://vnyiodyihbjosm4n.public.blob.vercel-storage.com/Tree-y3VjlYkkByWNSOol56jE307zHqaNsp.blend -O /tmp/Tree.blend",
+        shell=True,
+        check=True,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+    blend_bytes = Path("/tmp/Tree.blend").read_bytes()
+
+    Path(blend_path).write_bytes(blend_bytes)
+    Path(script_path).write_text(script_content)
+    Path(json_path).write_text(json_content)
+
+    cmd = ["blender", blend_path, "-b", "-P", script_path, "--", output_path, json_path]
+
+    subprocess.run(cmd, capture_output=True, text=True, cwd="/tmp")
+
+    output_file = Output(path=output_path)
+    output_file.save()
+    public_url = output_file.public_url(expires=400)
+    return {"output_url": public_url}
+
+
+if __name__ == "__main__":
+    output_url = render(
+        script_content=Path("render.py").read_text(),
+        json_content=Path("render_test.json").read_text(),
+        output_name="output.png",
+    )
+    print(f"Output URL: {output_url}")

blender/render.py

@@ -0,0 +1,63 @@
+import bpy
+import sys
+import json
+
+argv = sys.argv
+argv = argv[argv.index("--") + 1 :]
+
+output_path = argv[0]
+json_path = argv[1]
+
+with open(json_path, "r") as f:
+    config = json.load(f)
+
+
+def create_material(name, color):
+    if name in bpy.data.materials:
+        mat = bpy.data.materials[name]
+    else:
+        mat = bpy.data.materials.new(name=name)
+
+    mat.use_nodes = True
+    mat.node_tree.nodes.clear()
+
+    bsdf = mat.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled")
+    bsdf.inputs["Base Color"].default_value = (*color, 1.0)
+
+    output = mat.node_tree.nodes.new(type="ShaderNodeOutputMaterial")
+    mat.node_tree.links.new(bsdf.outputs["BSDF"], output.inputs["Surface"])
+
+    return mat
+
+
+if "material_color" in config:
+    color = config["material_color"]
+    new_material = create_material("TreeMaterial", color)
+
+    for obj in bpy.data.objects:
+        if obj.type == "MESH":
+            obj_name_lower = obj.name.lower()
+
+            if any(
+                keyword in obj_name_lower
+                for keyword in ["tree", "leaf", "leaves", "foliage"]
+            ):
+                obj.data.materials.clear()
+                obj.data.materials.append(new_material)
+
+if "object_to_modify" in config and "location" in config:
+    obj = bpy.data.objects.get(config["object_to_modify"])
+    if obj:
+        obj.location = config["location"]
+
+scene = bpy.context.scene
+scene.render.filepath = output_path
+scene.render.image_settings.file_format = "PNG"
+
+try:
+    scene.render.engine = "BLENDER_EEVEE"
+except:
+    scene.render.engine = "CYCLES"
+    scene.cycles.samples = 32
+
+bpy.ops.render.render(write_still=True)

blender/render_test.json

@@ -0,0 +1,5 @@
+{
+  "object_to_modify": "tree",
+  "location": [1.0, 2.0, 3.0],
+  "material_color": [0.1, 0.6, 0.1]
+}

image_generation/comfy_ui/controlnet/README.md

@@ -0,0 +1,45 @@
+# ComfyUI API Example
+
+An image generation API powered by ComfyUI, running ControlNet workflow
+
+## Deployment
+
+Deploy the app on Beam:
+
+```
+beam deploy app.py:handler
+```
+
+---
+
+## API Usage
+
+Send a `POST` request to the `/generate` endpoint with the following JSON body:
+
+```json
+{
+  "prompt": "Your prompt",
+  "image_url": "https://your-image-url"
+}
+```
+
+---
+
+### Example Request:
+
+```json
+{
+  "prompt": "A photorealistic golden retriever sitting in a field of flowers, soft light, professional lens, background blur",
+  "image_url": "https://encrypted-tbn1.gstatic.com/images?q=tbn:ANd9GcRHWNgtJvfMe4yIAOCPTpxRjPalxseEhGfMvW_B3Y3RLjhvwtCIHaYKt0D3K1h2qcWqqvvroakXVGKcAFnTDk7HhA"
+}
+```
+
+---
+
+### Example Response:
+
+```json
+{
+  "output_url": " https://app.beam.cloud/output/id/5cc90408-2c40-424f-bb3f-731268e7f100"
+}
+```

image_generation/comfy_ui/controlnet/controlnet.py

@@ -0,0 +1,168 @@
+from beam import Image, asgi, Output
+import requests
+from pathlib import Path
+
+image = (
+    Image()
+    .add_commands(["apt update && apt install git -y"])
+    .add_python_packages(
+        [
+            "fastapi[standard]==0.115.4",
+            "comfy-cli",
+            "huggingface_hub[hf_transfer]==0.26.2",
+        ]
+    )
+    .add_commands(
+        [
+            "yes | comfy install --nvidia --version 0.3.10",
+            "comfy node install was-node-suite-comfyui@1.0.2",
+            "mkdir -p /root/comfy/ComfyUI/models/checkpoints/",
+            "huggingface-cli download Comfy-Org/flux1-schnell flux1-schnell-fp8.safetensors --cache-dir /comfy-cache",
+            "ln -s /comfy-cache/models--Comfy-Org--flux1-schnell/snapshots/f2808ab17fe9ff81dcf89ed0301cf644c281be0a/flux1-schnell-fp8.safetensors /root/comfy/ComfyUI/models/checkpoints/flux1-schnell-fp8.safetensors",
+        ]
+    )
+)
+
+
+def download_image_from_url(url: str, filename: str = "input_image.png"):
+    target_path = Path("/root/comfy/ComfyUI/input") / filename
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    response = requests.get(url)
+    if response.status_code != 200:
+        raise Exception(f"Failed to download image: {response.status_code}")
+
+    with open(target_path, "wb") as f:
+        f.write(response.content)
+
+    return target_path
+
+
+def hf_download():
+    import subprocess
+    from huggingface_hub import hf_hub_download
+
+    dream_model = hf_hub_download(
+        repo_id="Bruhn/Lab_merge",
+        filename="dreamCreationVirtual3DECommerce_v10.safetensors",
+        cache_dir="/comfy-cache",
+    )
+
+    vae_model = hf_hub_download(
+        repo_id="stabilityai/sd-vae-ft-mse-original",
+        filename="vae-ft-mse-840000-ema-pruned.safetensors",
+        cache_dir="/comfy-cache",
+    )
+
+    controlnet_model = hf_hub_download(
+        repo_id="comfyanonymous/ControlNet-v1-1_fp16_safetensors",
+        filename="control_v11p_sd15_scribble_fp16.safetensors",
+        cache_dir="/comfy-cache",
+    )
+
+    subprocess.run(
+        "mkdir -p /root/comfy/ComfyUI/models/checkpoints/",
+        shell=True,
+        check=True,
+    )
+    subprocess.run(
+        "mkdir -p /root/comfy/ComfyUI/models/vae/",
+        shell=True,
+        check=True,
+    )
+    subprocess.run(
+        "mkdir -p /root/comfy/ComfyUI/models/controlnet/",
+        shell=True,
+        check=True,
+    )
+
+    subprocess.run(
+        f"ln -s {dream_model} /root/comfy/ComfyUI/models/checkpoints/dreamCreationVirtual3DECommerce_v10.safetensors",
+        shell=True,
+        check=True,
+    )
+    subprocess.run(
+        f"ln -s {vae_model} /root/comfy/ComfyUI/models/vae/vae-ft-mse-840000-ema-pruned.safetensors",
+        shell=True,
+        check=True,
+    )
+    subprocess.run(
+        f"ln -s {controlnet_model} /root/comfy/ComfyUI/models/controlnet/control_v11p_sd15_scribble_fp16.safetensors",
+        shell=True,
+        check=True,
+    )
+
+    cmd = "comfy launch --background"
+    subprocess.run(cmd, shell=True, check=True)
+
+
+@asgi(
+    name="comfy",
+    image=image,
+    on_start=hf_download,
+    cpu=8,
+    memory="32Gi",
+    gpu="A100-40",
+    timeout=-1,
+)
+def handler():
+    from fastapi import FastAPI, HTTPException
+    import subprocess
+    import json
+    from pathlib import Path
+    import uuid
+    from typing import Dict
+
+    app = FastAPI()
+
+    # This is where you specify the path to your workflow file.
+    # Make sure "workflow_api.json" exists in the same directory as this script.
+    WORKFLOW_FILE = Path(__file__).parent / "workflow_api.json"
+    OUTPUT_DIR = Path("/root/comfy/ComfyUI/output")
+
+    @app.post("/generate")
+    async def generate(item: Dict):
+        if not WORKFLOW_FILE.exists():
+            raise HTTPException(status_code=500, detail="Workflow file not found.")
+
+        image_url = item.get("image_url")
+        prompt = item.get("prompt")
+
+        if not image_url or not prompt:
+            raise HTTPException(status_code=400, detail="Missing image_url or prompt")
+
+        downloaded_image_path = download_image_from_url(image_url, "input.jpg")
+        request_id = uuid.uuid4().hex
+
+        workflow_data = json.loads(WORKFLOW_FILE.read_text())
+        workflow_data["6"]["inputs"]["text"] = item["prompt"]
+        workflow_data["11"]["inputs"]["image"] = str(downloaded_image_path)
+
+        new_workflow_file = Path(f"{request_id}.json")
+        new_workflow_file.write_text(json.dumps(workflow_data, indent=4))
+
+        # Run inference
+        cmd = (
+            f"comfy run --workflow {new_workflow_file} --wait --timeout 1200 --verbose"
+        )
+        subprocess.run(cmd, shell=True, check=True)
+
+        image_files = list(OUTPUT_DIR.glob("*"))
+
+        # Find the latest image
+        latest_image = max(
+            (f for f in image_files if f.suffix.lower() in {".png", ".jpg", ".jpeg"}),
+            key=lambda f: f.stat().st_mtime,
+            default=None,
+        )
+
+        if not latest_image:
+            raise HTTPException(status_code=404, detail="No output image found.")
+
+        output_file = Output(path=latest_image)
+        output_file.save()
+        public_url = output_file.public_url(expires=-1)
+        print(public_url)
+        return {"output_url": public_url}
+
+    return app