import gradio as gr
import bpy
import tempfile

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"])




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

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)
w