bpy-nvidia-T4 / app.py
kolibril13
add new script
4373daa
raw
history blame
2.93 kB
import gradio as gr
import bpy
from tqdm import tqdm
from math import pi
import tempfile
import molecularnodes as mn
import os
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) # <- this line would causes the error 🐛
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 = 1000
bpy.context.scene.render.resolution_x = 1000
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)