re-add molecular

app.py

@@ -1,6 +1,10 @@
 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"
@@ -22,21 +26,75 @@ def enable_GPUS():
         print(d["name"])
 
 
+# enable_GPUS()
 
+window = bpy.context.window
+screen = window.screen
 
-def generate():
-    with tempfile.NamedTemporaryFile(suffix=".JPEG", delete=False) as f:
-        bpy.context.scene.render.resolution_y = 200
-        bpy.context.scene.render.resolution_x = 400
-        bpy.context.scene.render.image_settings.file_format = "JPEG"
-        bpy.context.scene.render.filepath = f.name
-        enable_GPUS()
-        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
+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():
@@ -51,5 +109,4 @@ with gr.Blocks() as demo:
     )
 
 demo.queue(concurrency_count=1)
-demo.launch(debug=True, inline=True)
-w
+demo.launch(debug=True, inline=True)

requirements.txt

@@ -1 +1,2 @@
-bpy
+bpy
+molecularnodes