RayCasting-TRAVEL-ComfyUI-API.py

#This is an example that uses the websockets api to know when a prompt execution is done
#Once the prompt execution is done it downloads the images using the /history endpoint

import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid
import json
import urllib.request
import urllib.parse
import random

import pygame as pg
import cv2 as cv
import base64
from PIL import Image
from numba import njit
from threading import Thread
import numpy as np
import io
import math 
import multiprocessing

def encodeImage ( img, width=512, height=512 ):
    string_image = pg.image.tostring(img, 'RGB')
    temp_surf = pg.image.fromstring (string_image,(width, height),'RGB' )
    tmp_arr = pg.surfarray.array3d(temp_surf)            
    tmp_arr = cv.transpose(tmp_arr) # transpose or image is x/y-flipped
    retval, bytes = cv.imencode('.jpg', tmp_arr)
    encodedImage = base64.b64encode(bytes).decode('utf-8')        
    return encodedImage 
# encode the init images

cnetImgScreen = pg.Surface((512, 512))
cnetImg = pg.Surface((1024, 512))
lImg = pg.Surface((1024, 512))

scrCount = 1
theSeed = 1543888376

turnFactor = 0
walkSpeed = 0

saveIncremental = True # saves the images in project folder 
saveCnet = True
renderButtons = True

fImg = pg.image.load('init.jpg')
firstImage = encodeImage ( fImg , 512, 512 )
latentMask = encodeImage ( pg.image.load('mask.jpg'), 1024, 512 )
# # # lastImage = firstImage
lImg.blit(fImg, (0,0))
lImg.blit(fImg, (512,0))

objects = []  # buttons

server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())

def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req =  urllib.request.Request("http://{}/prompt".format(server_address), data=data)
    return json.loads(urllib.request.urlopen(req).read())

def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response:
        return response.read()

def get_history(prompt_id):
    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
        return json.loads(response.read())

def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break #Execution is done
        else:
            continue #previews are binary data

    history = get_history(prompt_id)[prompt_id]
    for o in history['outputs']:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
            output_images[node_id] = images_output

    return output_images



class Button():
    def __init__(self, x, y, width, height, key, buttonText='Button', onclickFunction=None, onePress=False):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.onclickFunction = onclickFunction
        self.onePress = onePress
        self.key = key

        self.fillColors = {
            'normal': '#ffffff',
            'hover': '#666666',
            'pressed': '#333333',
        }

        self.buttonSurface = pg.Surface((self.width, self.height))
        self.buttonRect = pg.Rect(self.x, self.y, self.width, self.height)

        font = pg.font.SysFont("Arial", 10)
        self.buttonSurf = font.render(buttonText, True, (20, 20, 20))

        self.alreadyPressed = False

        objects.append(self)

    def process(self, screen):

        mousePos = pg.mouse.get_pos()
        
        self.buttonSurface.blit(self.buttonSurf, [
            self.buttonRect.width/2 - self.buttonSurf.get_rect().width/2,
            self.buttonRect.height/2 - self.buttonSurf.get_rect().height/2
        ])
        screen.blit(self.buttonSurface, self.buttonRect)
        
        self.buttonSurface.fill(self.fillColors['normal'])
        if self.buttonRect.collidepoint(mousePos):
            self.buttonSurface.fill(self.fillColors['hover'])

            if pg.mouse.get_pressed(num_buttons=3)[0]:
                self.buttonSurface.fill(self.fillColors['pressed'])

                if self.onePress:
                    self.onclickFunction()
                   
                elif not self.alreadyPressed:
                    self.onclickFunction()
                    self.alreadyPressed = True
                return self.key

            else:
                self.alreadyPressed = False

# button callback functions that we don't really need 
def btnGo():
    print('-> go')#  
    
def btnLeft():
    print('-> left')
    
def btnRight():
    print('-> right')
    
def main():

    global objects
    global lastImage
    global firstImage
    global cnetImg
    global cnetImgScreen
    global renderButtons
    global renderFrame
    global lImg
    global turnFactor 
    global walkSpeed 
    
    size = 50    # size of the map
    posx, posy, posz = (1, 1, 0.5)
    rot, rot_v = (0, 0)
    lx, ly, lz = (size/2-0.5, size/2-0.5, 1)    
    mapc, maph, mapr = room_generator(posx, posy, size)
    res, res_o = 0, [64, 96, 112, 160, 192, 224, 448]
    width, height, mod, inc, sky, floor = adjust_resol(res_o[res])
   
    nuc = 8
    pool = multiprocessing.Pool(processes = nuc)
    renderFrame = True
    
    alphaSDImage = 0
    
    bench = []
    running = True

    pg.init()

    screen = pg.display.set_mode((1024, 512),pg.RESIZABLE)

    clock = pg.time.Clock()
    pg.mouse.set_visible(True)

    blitDelay = 0    
    
    goButton = Button(226, 480, 80, 40, 'go', 'Go', btnGo)
    leftButton = Button(126, 480, 80, 40, 'left', 'Turn', btnLeft)
    rightButton = Button(326, 480, 80, 40, 'right', 'Turn', btnRight)
    

    while running:
        
        for event in pg.event.get():
            if event.type == pg.QUIT:
                running = False
            if event.type == pg.KEYDOWN:
                if event.key == pg.K_ESCAPE:
                    running = False
                if event.key == ord('f'):              
                    alphaSDImage = 0 
                if event.key == ord('g'):              
                    alphaSDImage = 127
                if event.key == ord('h'):              
                    alphaSDImage = 255                    
                if event.key == ord('p'):              
                    renderFrame = True                        
                if event.key == ord('q'): # change resolution
                    if res > 0 :
                        res = res-1
                        width, height, mod, inc, sky, floor = adjust_resol(res_o[res])
                if event.key == ord('e'):
                    if res < len(res_o)-1 :
                        res = res+1
                        width, height, mod, inc, sky, floor = adjust_resol(res_o[res])


        param_values = []
        for j in range(height): #vertical loop 
            rot_j = rot_v + np.deg2rad(24 - j/mod)
            for i in range(width): #horizontal vision loop
                param_values.append([rot, i, j, inc, rot_j])
        tam = len(param_values)
        lista = []
        pixels = []
        
        for i in range(nuc):
            lista.append([i, param_values[i*int(tam/nuc):(i+1)*int(tam/nuc)],
                          mapc, maph, lx, ly, lz, mapr, posx, posy, posz, mod])

        retorno = pool.map(caster, lista)

        for i in range(nuc):
            pixels.append(retorno[i][1])

        pixels = np.reshape(pixels, (height,width,3))
        pixels = np.asarray(pixels)/np.sqrt(np.max(pixels))
             
                        
        # player's movement

        pressed_keys = pg.key.get_pressed()        
        posx, posy, rot, rot_v, walkSpeed, turnFactor = keyboardMovement(pressed_keys, posx, posy, rot, rot_v, maph, clock.tick()/500, walkSpeed, turnFactor)       

        rot = rot + turnFactor
        posx, posy = (posx + walkSpeed*np.cos(rot), posy + walkSpeed*np.sin(rot))

        # image        
        cnetImgScreen = pg.surfarray.make_surface((np.rot90(pixels*255)).astype('uint8'))
        cnetImgScreen = pg.transform.scale(cnetImgScreen, (512, 512))
        
        pg.draw.rect(cnetImg, (0,0,0), pg.Rect(0, 0, 1024, 512))      
        cnetImg.blit(cnetImgScreen, (512,0))
        
        # pg.Surface.set_alpha(lImg, alphaSDImage)

        screen.blit(lImg, (0,0) )
        screen.blit(cnetImgScreen, (512, 0) )
        
        if renderButtons == True:
            for object in objects:
                key = object.process(screen)
                
                if key == "go":                  
                    #et = 0.5 # 0.1 #clock.tick()/500
                    #posx, posy = (posx + et*np.cos(rot), posy + et*np.sin(rot))
                    #renderButtons = False
                    #renderFrame = True
                    walkSpeed = walkSpeed + 0.001
                    
                if key == "left":
                    turnFactor = turnFactor + 0.0001
                    #rot = rot + math.pi/4
                    #renderButtons = False
                    #renderFrame = True 
                if key == "right":
                    turnFactor = turnFactor - 0.0001
                    #rot = rot - math.pi/4
                    #renderButtons = False
                    #renderFrame = True  


        pg.display.flip()

        # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
        if renderFrame == True:
            blitDelay = blitDelay + 1
            delayTrigg = round(clock.get_fps()*2) # wait for next frame after blit
            if blitDelay >= delayTrigg:
                blitDelay = 0
                renderFrame = False # disable render until new image was loaded
                lImg.blit(cnetImg, (512,0) ) 
                encoded_image = encodeImage ( cnetImg, 1024, 512 ) # image for controlnet
                encoded_sd_image = encodeImage ( lImg, 1024, 512 ) # will be stored as global lastImage
                thread = Thread( target=startWorkflow, args=(encoded_image,cnetImg) )   #, args=(encoded_image,encoded_sd_image)
                thread.start()
                  
    stop_thread = True
    pg.quit()
    pool.close()

# run this in a thread or program will pause until response of SD
def startWorkflow( encodedImage,cnetImg ):

    global renderFrame
    global renderButtons
    global cnetImgScreen

    # Opening JSON workflow
    # you can edit it while the script is running, if you like
    f = open('workflow_api.json')
     
    # returns JSON object as 
    # a dictionary
    prompt = json.load(f)

    # override the text prompt for our positive CLIPTextEncode
    #prompt["6"]["inputs"]["text"] = "masterpiece best quality man"

    # override the seed for our KSampler node
    # prompt["3"]["inputs"]["seed"] = 704184315842712 #random.random() * 1000000000
    
    # you can send the image through the A1111/SDNext API as a binary blob
    # but this doesn't seem to work (yet), so I...
    # save current depth map as temporary file (after I created the folder inside the /input folder) 
    pg.image.save(cnetImgScreen, "E:\\AI\\automatic-fork\\ComfyUI_windows_portable\\ComfyUI\\input\\API-IN\\depthMap_tmp.jpg")
    
    # point prompt to temporary depth map
    # the index may change if you edit the workflow.json! 
    prompt["69"]["inputs"]["image"] = "API-IN\\depthMap_tmp.jpg"

    ws = websocket.WebSocket()
    ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
    images = get_images(ws, prompt)

    for node_id in images:
        for image_data in images[node_id]:
            lImgTmp = pg.image.load( io.BytesIO(image_data) )
            lImg.blit(lImgTmp, (0,0) )
    
    renderButtons = True
    renderFrame = True

def room_generator(posx, posy, size):
    mapc = np.random.uniform(0, 0.1, (size,size,3)) 
    mapr = np.random.choice([0, 0, 0, 0], (size,size))
    maph = np.random.choice([0, 1], (size,size))
    maph[0,:], maph[size-1,:], maph[:,0], maph[:,size-1] = (1,1,1,1)

    mapc[posx][posy], maph[posx][posy], mapr[posx][posy] = (0, 0, 0)
    
    x,y = (0,0)
    
    for x in range(size):
        if x%2 == 0:
            for y in range(size):
                mapc[x][y], maph[x][y], mapr[x][y] = (0, 0, 0)

    for y in range(size):
        if y%2 == 0:
            for x in range(size):
                mapc[x][y], maph[x][y], mapr[x][y] = (0, 0, 0)

                
    return mapc, maph, mapr
    
def maze_generator(x, y, size):
    mapc = np.random.uniform(0, 0.1, (size,size,3)) 
    mapr = np.random.choice([0, 0, 0, 0], (size,size))
    maph = np.random.choice([0, 0, 0, 0, 0, 0, 0, 0, 1], (size,size))
    maph[0,:], maph[size-1,:], maph[:,0], maph[:,size-1] = (1,1,1,1)

    mapc[x][y], maph[x][y], mapr[x][y] = (0, 0, 0)
    count = 0 
    
    while 1:
        testx, testy = (x, y)
        
        if np.random.uniform() > 0.5:
            testx = testx + np.random.choice([-1, 1])
        else:
            testy = testy + np.random.choice([-1, 1])
            
        if testx > 0 and testx < size -1 and testy > 0 and testy < size -1:
        
            if maph[testx][testy] == 0 or count > 5:
                count = 0
                x, y = (testx, testy)
                mapc[x][y], maph[x][y], mapr[x][y] = (0, 0, 0)
                if x == size-2:
                    exitx, exity = (x, y)
                    break
            else:
                count = count+1
                
    return mapc, maph, mapr, exitx, exity

def keyboardMovement(pressed_keys,posx, posy, rot, rot_v, maph, et, walkSpeed, turnFactor):
    
    x, y = (posx, posy)
    
    # p_mouse = pg.mouse.get_pos()
    # rot = rot + 4*np.pi*(0.5-(p_mouse[0]-400)/150000)
    # rot_v = np.cos(et*0.1) * 0.1 # vertical view angle
    
    if pressed_keys[pg.K_UP] or pressed_keys[ord('w')]:
        x, y = (x + et*np.cos(rot), y + et*np.sin(rot))
        
    if pressed_keys[pg.K_DOWN] or pressed_keys[ord('s')]:
        #x, y = (x - et*np.cos(rot), y - et*np.sin(rot))
        walkSpeed = 0
        turnFactor = 0
        
    if pressed_keys[pg.K_LEFT] or pressed_keys[ord('a')]:
        # x, y = (x - et*np.sin(rot), y + et*np.cos(rot))
        rot = rot + 0.1
        
    if pressed_keys[pg.K_RIGHT] or pressed_keys[ord('d')]:
        #x, y = (x + et*np.sin(rot), y - et*np.cos(rot))
        rot = rot - 0.1

    if maph[int(x)][int(y)] == 0:
        posx, posy = (x, y)
        
    
                                                
    return posx, posy, rot, rot_v, walkSpeed, turnFactor
        
       
@njit(fastmath=True)
def fast_ray(x, y, z, cos, sin, sinz, maph):
    while 1:
        x, y, z = x + cos, y + sin, z + sinz
        if (z > 1 or z < 0):
            break
        if maph[int(x)][int(y)] > z:
            break        
    return x, y, z        

def view_ray(x, y, z, cos, sin, sinz, mapc, lx, ly, lz, maph):
    
    x, y, z = fast_ray(x, y, z, cos, sin, sinz, maph)
    dtol = np.sqrt((x-lx)**2+(y-ly)**2+(lz-1)**2)

    if z > 1: # ceiling
##        c = np.asarray([0.3,0.7,1])
        c = np.asarray([1,1,1]) 
    elif z < -1: # floor
        c = np.asarray([.3,.3,.3])
    else:
        c = np.asarray([1,1,1]) # if all fails

    h = 1*np.clip(1/dtol, 0, 1)
    c = c*h
    return c, x, y, z, dtol

@njit(fastmath=True)
def shadow_ray(x, y, z, lx, ly, lz, maph, c, inc, dtol):
    dx, dy, dz = inc*5*(lx-x)/dtol, inc*5*(ly-y)/dtol, inc*5*(lz-z)/dtol
    mod = 1
    while 1:
        x, y, z = (x + dx, y + dy, z + dz)
        if maph[int(x)][int(y)]!= 0 and z<= maph[int(x)][int(y)]:
            mod = mod*0.9
            if mod < 0.5:
                break
        elif z > 0.9:
            break
    return c*mod

def reflection(x, y, z, cos, sin, sinz, mapc, lx, ly, lz, maph, exitx, exity, c, posz, inc, mapr, recur):
    if abs(z-maph[int(x)][int(y)])<abs(sinz):
        sinz = -sinz
    elif maph[int(x+cos)][int(y-sin)] != 0:
        cos = -cos
    else:
        sin = -sin
    c2, x, y, z, dtol = view_ray(x, y, z, cos, sin, sinz, mapc, lx, ly, lz, maph, exitx, exity)
    if z < 1:
        c2 = shadow_ray(x, y, z, lx, ly, lz, maph, c2, inc, dtol)
    if (mapr[int(x)][int(y)] != 0 and z < 1 and z > 0 and not recur):
        c2 = reflection(x, y, z, cos, sin, sinz, mapc, lx, ly, lz, maph, exitx, exity, c2, posz, inc, mapr, recur=True)
    c = (c + c2)/2
    return c

def caster(lista):
    param_values = lista[1]
    mapc = lista[2]
    maph = lista[3]
    lx = lista[8] # put the light on player position
    ly = lista[9]
    lz = lista[10] - 1
    #exitx = lista[7]
    #exity = lista[8]
    
    mapr = lista[7]
    posx = lista[8]
    posy = lista[9]
    posz = lista[10]
    mod = lista[11]
    
    pixels = []
    
    for values in param_values:
        rot = values[0]
        i = values[1]
        j = values[2]
        inc = values[3]
        rot_j = values[4]
        rot_i = rot + np.deg2rad(i/mod - 30)
        x, y, z = (posx, posy, posz)
        sin, cos,  = (inc*np.sin(rot_i), inc*np.cos(rot_i))
        sinz = inc*np.sin(rot_j)
        c, x, y, z, dtol = view_ray(x, y, z, cos, sin, sinz, mapc, lx, ly, lz,
                                    maph)
        #if z < 1:
            #c = shadow_ray(x, y, z, lx, ly, lz, maph, c, inc, dtol)
            #if mapr[int(x)][int(y)] != 0 and z > 0:
            #    c = reflection(x, y, z, cos, sin, sinz, mapc, lx, ly, lz, maph,
            #                   exitx, exity, c, posz, inc, mapr, recur=False)
                    
        pixels.append(c)

    return lista[0], pixels

def ray_caster(x, y, i, ex, ey, maph, mapc, sin, cos, n, half, mod):
    zz= 0.5
    if half == None:
        zz = 0.1
    x, y, n, tc, ty = fast_ray_caster(x, y, zz, cos, sin, maph, n, i, ex, ey, mod)
    h , c = shader(n, maph, mapc, sin, cos, x, y, i, mod)
    
    if maph[int(x)][int(y)] < 0.5 and half == None:
        half = [h, c, n]
        x, y, n, tc2, ty2 = fast_ray_caster(x, y, 0.5, cos, sin, maph, n, i, ex, ey, mod)
        ty, tc = ty + ty2, tc + tc2
        h , c = shader(n, maph, mapc, sin, cos, x, y, i, mod)
    return(c, h, x, y, n, half, ty, tc)


@njit(fastmath=True)
def fast_ray_caster(x, y, z, cos, sin, maph, n, i, ex, ey, mod):
    ty, tc = [], []
    while 1:
        n = n+1
        x, y = x + cos, y + sin
        if z < 0.5 and int(x*2)%2 == int(y*2)%2:
                th = 1/(0.05/mod * n)#*np.cos(np.deg2rad(i/mod - 30)))
                if th < 1  and th >= 0:
                    ty.append(th)
                    if int(x) == ex and int(y) == ey:
                        tc.append(np.asarray([0,0,1]))
                    else:
                        tc.append(np.asarray([0,0,0]))
        if maph[int(x)][int(y)] > z:
            break        
    return x, y, n, tc, ty

def shader(n, maph, mapc, sin, cos, x, y, i, mod):
    
    h = np.clip(1/(0.05/mod * n), 0, 1)#*np.cos(np.deg2rad(i/mod-30))), 0, 1)
    c = np.asarray(mapc[int(x)][int(y)])*(0.4 + 0.6 * h)
    
    if maph[int(x+cos)][int(y-sin)] == 1:
        c = 0.85*c
        
        if maph[int(x-cos)][int(y+sin)] == 1 and sin >0:
            c = 0.7*c
    return h, c

def reflection_caster(x, y, i, ex, ey, maph, mapc, sin, cos, n, c, h, half, pixels, ty, tc, height, mod):
    
    hor = int(height/2)
    hh = int((h*height)/2)
    pixels[hor-hh:hor+hh,i] = np.add(pixels[hor-hh:hor+hh,i], np.asarray([c]*(hh*2)))/2
    
    if maph[int(x+cos)][int(y-sin)] > 0.5:
        cos = -cos
        
    else:
        sin = -sin
        
    c2, h2, x, y, n2, half2, ty2, tc2 = ray_caster(x, y, i, ex, ey, maph, mapc, sin, cos, n, half, mod)
        
    ty, tc = ty + ty2, tc + tc2
    hh = int((h2*height)/2)
    pixels[hor-hh:hor+hh,i] = (c + c2)/2
    
    if half2 != None and half == None:
        hh = int((half2[0]*height)/2)
        pixels[hor:hor+hh,i] = (c + half2[1])/2
        
    elif half != None:
        hh = int((half[0]*height)/2)
        pixels[hor:hor+hh,i] = half[1]
           
    return pixels, ty, tc     

def adjust_resol(width):
    height = width
    mod = width/64
    inc = 0.05/mod
    gradient = np.linspace(0,1,int(height/2-1))
    sky = np.asarray([gradient,gradient,gradient]).T
    floor = np.asarray([gradient,gradient,gradient]).T
    print('Resolution: ', width, height)
    return width, height, mod, inc, sky, floor

if __name__ == '__main__':
    main()