Spaces:
Runtime error
Runtime error
File size: 4,768 Bytes
9f012c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import os
import asyncio
import gradio as gr
from PIL import Image
from dotenv import load_dotenv, find_dotenv
from naptha_sdk.client.naptha import Naptha
from naptha_sdk.user import generate_user
from naptha_sdk.schemas import ModuleRun
load_dotenv(find_dotenv())
async def run_module(module_name, prompt, selected_node):
user, _ = generate_user(os.getenv("PRIVATE_KEY"))
hub_url = os.getenv("HUB_URL")
hub_username = os.getenv("HUB_USER")
hub_password = os.getenv("HUB_PASS")
# Use the selected node
if selected_node == "node":
node_url = 'http://node.naptha.ai:7001'
elif selected_node == "node1":
node_url = 'http://node1.naptha.ai:7001'
else:
return None, f"Error: Invalid node selection: {selected_node}"
routing_url = os.getenv("ROUTING_URL", None)
indirect_node_id = os.getenv("INDIRECT_NODE_ID", None)
output_dir = './naptha_output'
if module_name == 'generate_image':
os.makedirs(output_dir, exist_ok=True)
for f in os.listdir(output_dir):
os.remove(os.path.join(output_dir, f))
naptha = await Naptha(
user=user,
hub_username=hub_username,
hub_password=hub_password,
hub_url=hub_url,
node_url=node_url,
routing_url=routing_url,
indirect_node_id=indirect_node_id,
)
module_run_input = {
'consumer_id': naptha.user["id"],
"module_name": module_name,
'worker_nodes': None,
"module_params": {
"prompt": prompt
},
}
print(f"Running module {module_name} on {selected_node} with parameters: {module_run_input}")
print("Checking user...")
user = await naptha.node.check_user(user_input=naptha.user)
if user["is_registered"]:
print("Found user...", user)
else:
print("No user found. Registering user...")
user = await naptha.node.register_user(user_input=user)
print(f"User registered: {user}.")
print("Running...")
module_run = await naptha.node.run_task(module_run_input)
if isinstance(module_run, dict):
module_run = ModuleRun(**module_run)
print(f"Module Run ID: {module_run.id}")
current_results_len = 0
while True:
module_run = await naptha.node.check_task(module_run)
if isinstance(module_run, dict):
module_run = ModuleRun(**module_run)
output = f"{module_run.status} {module_run.module_type} {module_run.module_name}"
if len(module_run.child_runs) > 0:
output += f", task {len(module_run.child_runs)} {module_run.child_runs[-1].module_name} (node: {module_run.child_runs[-1].worker_nodes[0]})"
print(output)
if len(module_run.results) > current_results_len:
print("Output: ", module_run.results[-1])
current_results_len += 1
if module_run.status in ['completed', 'error']:
break
await asyncio.sleep(3)
if module_run.status == 'completed':
results = module_run.results
print(f"Results: {results}")
if module_name == 'generate_image':
folder_id = module_run.id.split(":")[1]
image_path = f"{output_dir}/output.png"
await naptha.node.read_storage(folder_id.strip(), output_dir, ipfs=None)
if os.path.exists(image_path):
return image_path, None
else:
return None, "Image generation failed or image not found."
else: # chat
return None, results[-1] if results else "No output generated."
else:
error_message = module_run.error_message
print(f"Error: {error_message}")
return None, f"Error: {error_message}"
def gradio_wrapper(module, prompt, node):
image_path, text = asyncio.run(run_module(module, prompt, node))
if image_path:
return image_path, gr.update(visible=False)
elif text:
return gr.update(visible=False), text
else:
return gr.update(visible=False), gr.update(visible=False)
with gr.Blocks() as iface:
gr.Markdown("# Naptha Module Runner")
gr.Markdown("Run either the 'generate_image' or 'chat' module using Naptha on selected node.")
with gr.Row():
module = gr.Dropdown(["generate_image", "chat"], label="Module")
node = gr.Dropdown(["node", "node1"], label="Node")
prompt = gr.Textbox(label="Prompt")
with gr.Row():
image_output = gr.Image(label="Generated Image")
text_output = gr.Textbox(label="Chat Output")
submit_btn = gr.Button("Run")
submit_btn.click(
fn=gradio_wrapper,
inputs=[module, prompt, node],
outputs=[image_output, text_output]
)
iface.launch() |