Spaces:
Sleeping
Sleeping
File size: 3,597 Bytes
4373daa 05f01a3 4373daa 05f01a3 4373daa 6fa48ff f3f1ddc 6fa48ff 05f01a3 9108d7c 05f01a3 9108d7c 4373daa 05f01a3 |
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 |
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) |