Spaces:
Running
Running
| 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) |