File size: 3,642 Bytes
4373daa
 
 
 
 
 
 
 
9108d7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4373daa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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) # <- 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)