kolibril13
add gpu
f3f1ddc
raw
history blame contribute delete
No virus
3.6 kB
import gradio as gr
import bpy
from tqdm import tqdm
from math import pi
import tempfile
import molecularnodes as mn
import os
def enable_GPUS():
bpy.data.scenes[0].render.engine = "CYCLES" #"CYCLES"
# Set the device_type
bpy.context.preferences.addons[
"cycles"
].preferences.compute_device_type = "CUDA" # or "OPENCL"
# Set the device and feature set
bpy.context.scene.cycles.device = "GPU"
for scene in bpy.data.scenes:
scene.cycles.device = "GPU"
bpy.context.preferences.addons["cycles"].preferences.get_devices()
print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type)
for d in bpy.context.preferences.addons["cycles"].preferences.devices:
d["use"] = True # Using all devices, include GPU and CPU
print(d["name"])
enable_GPUS()
window = bpy.context.window
screen = window.screen
style = 'cartoon'
nodes_to_append = ["MN_color_set",
"MN_color_common",
"MN_color_attribute_random",
mn.nodes.styles_mapping[style]]
def get_areas(type):
return [area for area in screen.areas if area.type == type]
def get_regions(areas):
return [region for region in areas[0].regions if region.type == 'WINDOW']
for node in nodes_to_append:
bpy.ops.wm.append(
'INVOKE_DEFAULT',
directory = os.path.join(mn.nodes.mn_data_file, 'NodeTree'),
filename = node,
link = False
)
def generate(progress=gr.Progress(track_tqdm=True)):
area_type = 'VIEW_3D'
areas = get_areas(area_type)
with bpy.context.temp_override(window=window, area=areas[0], region=get_regions(areas)[0], screen=screen):
for obj in bpy.context.scene.objects:
if obj.type == 'MESH':
bpy.data.objects.remove(obj, do_unlink=True)
molecule = mn.load.molecule_rcsb("7TYG", starting_style=style, center_molecule=True)
molecule.select_set(True)
bpy.context.view_layer.objects.active = molecule
bpy.ops.view3d.camera_to_view_selected()
camera = bpy.data.objects["Camera"]
camera.data.dof.use_dof = True
camera.data.dof.focus_distance = 5
camera.data.dof.aperture_fstop = 4
camera.data.angle = pi / 3
camera.data.type = "PERSP"
with tempfile.NamedTemporaryFile(suffix=".JPEG", delete=False) as f:
bpy.context.scene.render.resolution_y = 400
bpy.context.scene.render.resolution_x = 600
bpy.context.scene.render.image_settings.file_format = "JPEG"
bpy.context.scene.render.filepath = f.name
with tqdm() as pbar:
def elapsed(dummy):
pbar.update()
bpy.app.handlers.render_stats.append(elapsed)
bpy.context.scene.frame_set(1)
bpy.context.scene.frame_current = 1
bpy.ops.render.render(animation=False, write_still=True)
bpy.data.images["Render Result"].save_render(
filepath=bpy.context.scene.render.filepath
)
bpy.app.handlers.render_stats.clear()
return f.name
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
render_btn = gr.Button("Render")
with gr.Column(scale=3):
image = gr.Image(type="filepath")
render_btn.click(
generate,
outputs=[image],
)
demo.queue(concurrency_count=1)
demo.launch(debug=True, inline=True)