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)