preferences.py

import bpy
from bpy.props import CollectionProperty, StringProperty
from bpy_extras.io_utils import ImportHelper
import os
import webbrowser

from .absolute_path import absolute_path
from .operators.install_dependencies import InstallDependencies
from .operators.open_latest_version import OpenLatestVersion
from .ui.presets import RestoreDefaultPresets, default_presets_missing
from .generator_process import Generator
from .generator_process.actions.prompt_to_image import Pipeline
from .generator_process.actions.huggingface_hub import DownloadStatus, ModelType
from .generator_process.actions.convert_original_stable_diffusion_to_diffusers import ModelConfig

is_downloading = False

class OpenHuggingFace(bpy.types.Operator):
    bl_idname = "dream_textures.open_hugging_face"
    bl_label = "Get Access Token"
    bl_description = ("Opens huggingface.co to the tokens page")
    bl_options = {"REGISTER", "INTERNAL"}

    def execute(self, context):
        webbrowser.open("https://huggingface.co/settings/tokens")
        return {"FINISHED"}

_model_config_options = [(m.name, m.value, '') for m in ModelConfig]
class ImportWeights(bpy.types.Operator, ImportHelper):
    bl_idname = "dream_textures.import_weights"
    bl_label = "Import Checkpoint File"
    filename_ext = ".ckpt"
    filter_glob: bpy.props.StringProperty(
        default="*.ckpt",
        options={'HIDDEN'},
        maxlen=255,
    )
    model_config: bpy.props.EnumProperty(
        name="Model Config",
        items=_model_config_options
    )

    def execute(self, context):
        _, extension = os.path.splitext(self.filepath)
        if extension != '.ckpt':
            self.report({"ERROR"}, "Select a valid stable diffusion '.ckpt' file.")
            return {"FINISHED"}
        try:
            Generator.shared().convert_original_stable_diffusion_to_diffusers(self.filepath, ModelConfig[self.model_config]).result()
        except Exception as e:
            self.report({"ERROR"}, str(e))
        
        set_model_list('installed_models', Generator.shared().hf_list_installed_models().result())

        return {"FINISHED"}

class OpenContributors(bpy.types.Operator):
    bl_idname = "dream_textures.open_contributors"
    bl_label = "See All Contributors"

    def execute(self, context):
        webbrowser.open("https://github.com/carson-katri/dream-textures/graphs/contributors")
        return {"FINISHED"}

class OpenDreamStudio(bpy.types.Operator):
    bl_idname = "dream_textures.open_dream_studio"
    bl_label = "Find Your Key"
    bl_description = ("Opens DreamStudio to the API key tab.")
    bl_options = {"REGISTER", "INTERNAL"}

    def execute(self, context):
        webbrowser.open("https://beta.dreamstudio.ai/membership?tab=apiKeys")
        return {"FINISHED"}

class Model(bpy.types.PropertyGroup):
    bl_label = "Model"
    bl_idname = "dream_textures.Model"

    model: bpy.props.StringProperty(name="Model")
    downloads: bpy.props.IntProperty(name="Downloads")
    likes: bpy.props.IntProperty(name="Likes")
    model_type: bpy.props.EnumProperty(name="Model Type", items=[(t.name, t.name, '') for t in ModelType])

class PREFERENCES_UL_ModelList(bpy.types.UIList):
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname):
        model_name = item.model
        is_installed = False
        if os.path.exists(item.model):
            model_name = os.path.basename(item.model).replace('models--', '').replace('--', '/')
            is_installed = True
        split = layout.split(factor=0.75)
        split.label(text=model_name)
        if item.downloads != -1:
            split.label(text=str(item.downloads), icon="IMPORT")
        if item.downloads != -1:
            split.label(text=str(item.likes), icon="HEART")
        if ModelType[item.model_type] != ModelType.UNKNOWN:
            split.label(text=item.model_type.replace('_', ' ').title())
        layout.operator(InstallModel.bl_idname, text="", icon="FILE_FOLDER" if is_installed else "IMPORT").model = item.model

@staticmethod
def set_model_list(model_list: str, models: list):
    getattr(bpy.context.preferences.addons[__package__].preferences, model_list).clear()
    for model in models:
        m = getattr(bpy.context.preferences.addons[__package__].preferences, model_list).add()
        m.model = model.id
        m.downloads = model.downloads
        m.likes = model.likes
        try:
            m.model_type = model.model_type.name
        except:
            pass

class ModelSearch(bpy.types.Operator):
    bl_idname = "dream_textures.model_search"
    bl_label = "Search"
    bl_description = ("Searches Hugging Face Hub for models")
    bl_options = {"REGISTER", "INTERNAL"}

    def execute(self, context):
        
        return {"FINISHED"}

class InstallModel(bpy.types.Operator):
    bl_idname = "dream_textures.install_model"
    bl_label = "Install or Open"
    bl_description = ("Install or open a model from the cache")
    bl_options = {"REGISTER", "INTERNAL"}

    model: StringProperty(name="Model ID")

    def execute(self, context):
        if os.path.exists(self.model):
            webbrowser.open(f"file://{self.model}")
        else:
            global is_downloading
            is_downloading = True
            f = Generator.shared().hf_snapshot_download(self.model, bpy.context.preferences.addons[__package__].preferences.hf_token)
            def on_progress(_, response: DownloadStatus):
                bpy.context.preferences.addons[__package__].preferences.download_file = response.file
                bpy.context.preferences.addons[__package__].preferences.download_progress = int((response.index / response.total) * 100)
            def on_done(future):
                global is_downloading
                is_downloading = False
                set_model_list('installed_models', Generator.shared().hf_list_installed_models().result())
            def on_exception(_, exception):
                self.report({"ERROR"}, str(exception))
            f.add_response_callback(on_progress)
            f.add_done_callback(on_done)
            f.add_exception_callback(on_exception)
        return {"FINISHED"}

def _model_search(self, context):
    def on_done(future):
        set_model_list('model_results', future.result())
    Generator.shared().hf_list_models(self.model_query).add_done_callback(on_done)

def _update_ui(self, context):
    if hasattr(context.area, "regions"):
        for region in context.area.regions:
            if region.type == "UI":
                region.tag_redraw()
    return None

class StableDiffusionPreferences(bpy.types.AddonPreferences):
    bl_idname = __package__

    dream_studio_key: StringProperty(name="DreamStudio Key")

    model_query: StringProperty(name="Search", update=_model_search)
    model_results: CollectionProperty(type=Model)
    active_model_result: bpy.props.IntProperty(name="Active Model", default=0)
    hf_token: StringProperty(name="HuggingFace Token")

    installed_models: CollectionProperty(type=Model)
    active_installed_model: bpy.props.IntProperty(name="Active Model", default=0)

    download_file: bpy.props.StringProperty(name="")
    download_progress: bpy.props.IntProperty(name="", min=0, max=100, subtype="PERCENTAGE", update=_update_ui)

    @staticmethod
    def register():
        if Pipeline.local_available():
            def on_done(future):
                set_model_list('installed_models', future.result())
            Generator.shared().hf_list_installed_models().add_done_callback(on_done)

    def draw(self, context):
        layout = self.layout

        weights_installed = len(self.installed_models) > 0

        if not weights_installed:
            layout.label(text="Complete the following steps to finish setting up the addon:")

        has_dependencies = len(os.listdir(absolute_path(".python_dependencies"))) > 2
        if has_dependencies:
            has_local = Pipeline.local_available()

            if has_local:
                if is_downloading:
                    progress_col = layout.column()
                    progress_col.label(text=f"Downloading {self.download_file}")
                    progress_col.prop(self, "download_progress", slider=True)
                    progress_col.enabled = False
                else:
                    if not weights_installed:
                        default_weights_box = layout.box()
                        default_weights_box.label(text="You need to download at least one model.")
                        default_weights_box.operator(InstallModel.bl_idname, text="Download Stable Diffusion v2.1 (Recommended)", icon="IMPORT").model = "stabilityai/stable-diffusion-2-1"

                    search_box = layout.box()
                    search_box.label(text="Find Models", icon="SETTINGS")
                    search_box.label(text="Search Hugging Face Hub for more compatible models.")

                    search_box.prop(self, "model_query", text="", icon="VIEWZOOM")
                    
                    if len(self.model_results) > 0:
                        search_box.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_model_results", self, "model_results", self, "active_model_result")

                    search_box.label(text="Some models require authentication. Provide a token to download gated models.")

                    auth_row = search_box.row()
                    auth_row.prop(self, "hf_token", text="Token")
                    auth_row.operator(OpenHuggingFace.bl_idname, text="Get Your Token", icon="KEYINGSET")

                layout.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_installed_models", self, "installed_models", self, "active_installed_model")
                layout.operator(ImportWeights.bl_idname, icon='IMPORT')
            
            dream_studio_box = layout.box()
            dream_studio_box.label(text=f"DreamStudio{' (Optional)' if has_local else ''}", icon="HIDE_OFF")
            dream_studio_box.label(text=f"Link to your DreamStudio account to run in the cloud{' instead of locally.' if has_local else '.'}")
            key_row = dream_studio_box.row()
            key_row.prop(self, "dream_studio_key", text="Key")
            key_row.operator(OpenDreamStudio.bl_idname, text="Find Your Key", icon="KEYINGSET")

            if weights_installed or len(self.dream_studio_key) > 0:
                complete_box = layout.box()
                complete_box.label(text="Addon Setup Complete", icon="CHECKMARK")
                complete_box.label(text="To locate the interface:")
                complete_box.label(text="1. Open an Image Editor or Shader Editor space")
                complete_box.label(text="2. Enable 'View' > 'Sidebar'")
                complete_box.label(text="3. Select the 'Dream' tab")
            
            if default_presets_missing():
                presets_box = layout.box()
                presets_box.label(text="Default Presets", icon="PRESET")
                presets_box.label(text="It looks like you removed some of the default presets.")
                presets_box.label(text="You can restore them here.")
                presets_box.operator(RestoreDefaultPresets.bl_idname, icon="RECOVER_LAST")
        else:
            missing_dependencies_box = layout.box()
            missing_dependencies_box.label(text="Dependencies Missing", icon="ERROR")
            missing_dependencies_box.label(text="You've likely downloaded source instead of release by accident.")
            missing_dependencies_box.label(text="Follow the instructions to install for your platform.")
            missing_dependencies_box.operator(OpenLatestVersion.bl_idname, text="Download Latest Release")
        
        contributors_box = layout.box()
        contributors_box.label(text="Contributors", icon="COMMUNITY")
        contributors_box.label(text="Dream Textures is made possible by the contributors on GitHub.")
        contributors_box.operator(OpenContributors.bl_idname, icon="URL")

        if context.preferences.view.show_developer_ui: # If 'Developer Extras' is enabled, show addon development tools
            developer_box = layout.box()
            developer_box.label(text="Development Tools", icon="CONSOLE")
            developer_box.label(text="This section is for addon development only. You are seeing this because you have 'Developer Extras' enabled.")
            developer_box.label(text="Do not use any operators in this section unless you are setting up a development environment.")
            if has_dependencies:
                warn_box = developer_box.box()
                warn_box.label(text="Dependencies already installed. Only install below if you developing the addon", icon="CHECKMARK")
            developer_box.prop(context.scene, 'dream_textures_requirements_path')
            developer_box.operator(InstallDependencies.bl_idname, icon="CONSOLE")