app.py

import subprocess
import os
import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageEnhance
from pathlib import Path

# Constants
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = np.iinfo(np.int32).max
OUTPUT_DIR = Path("output_minecraft_skins")

print(f"Using: {DEVICE}")

# Repo URL
REPO_URL = "https://github.com/BF667/Minecraft_Skin_Generator.git"
REPO_NAME = "Minecraft_Skin_Generator"
REPO_PATH = Path(REPO_NAME)

def setup_repository():
    """Clones or updates the Git repository."""
    if not REPO_PATH.exists():
        print(f"Cloning {REPO_NAME} repository...")
        try:
            subprocess.run(["git", "clone", REPO_URL], check=True)
            print("Repository cloned successfully.")
        except subprocess.CalledProcessError as e:
            print(f"Error cloning repository: {e}")
            raise
    else:
        print(f"{REPO_NAME} repository already exists. Checking for updates...")
        try:
            # Change to the repository directory to pull updates
            os.chdir(REPO_PATH)
            subprocess.run(["git", "pull"], check=True)
            print("Repository updated successfully.")
            os.chdir("..") # Change back to original directory
        except subprocess.CalledProcessError as e:
            print(f"Error updating repository: {e}")
            raise

    # Change to the repository directory for script execution
    os.chdir(REPO_PATH)


def run_inference(
    prompt: str,
    stable_diffusion_model: str,
    num_inference_steps: int,
    guidance_scale: float,
    model_precision_type: str,
    seed: int,
    filename: str,
    model_3d: bool,
    verbose: bool,
):
    """Runs the inference process for generating Minecraft skins."""

    sd_model_map = {
        '2': "minecraft-skins",
        'xl': "minecraft-skins-sdxl",
    }
    sd_script = sd_model_map.get(stable_diffusion_model)

    if not sd_script:
        raise ValueError(f"Invalid stable_diffusion_model: {stable_diffusion_model}")

    command = [
        "python",
        f"Scripts/{sd_script}.py",
        prompt,
        str(num_inference_steps),
        str(guidance_scale),
        model_precision_type,
        str(seed),
        filename,
    ]

    if model_3d:
        command.append("--model_3d")
    if verbose:
        command.append("--verbose")

    try:
        # Use subprocess.run for better control and error handling
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        print("Inference command output:")
        print(result.stdout)
        if result.stderr:
            print("Inference command error output:")
            print(result.stderr)

    except subprocess.CalledProcessError as e:
        print(f"Error during inference: {e}")
        print(f"Stdout: {e.stdout}")
        print(f"Stderr: {e.stderr}")
        return None, None  # Return None for outputs on error
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        return None, None

    # Ensure output directory exists
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Construct output paths using pathlib
    image_path = OUTPUT_DIR / filename
    model_3d_path = OUTPUT_DIR / f"{filename.split('.')[0]}_3d_model.glb" if model_3d else None

    # Basic check for file existence (can be more robust)
    if not image_path.exists():
        print(f"Warning: Image file not found at {image_path}")
        image_path = None
    if model_3d and model_3d_path and not model_3d_path.exists():
        print(f"Warning: 3D model file not found at {model_3d_path}")
        model_3d_path = None

    return str(image_path) if image_path else None, str(model_3d_path) if model_3d_path else None


def create_gradio_ui():
    """Defines and returns the Gradio UI components."""
    with gr.Blocks(title="Minecraft Skin Generator", css=".pixelated {image-rendering: pixelated} .checkered img {background-image: url(\'data:image/svg+xml,<svg xmlns=\'http://www.w3.org/2000/svg\' width=\'2\' height=\'2\' fill-opacity=\'.15\'><rect x=\'1\' width=\'1\' height=\'1\'/><rect y=\'1\' width=\'1\' height=\'1\'/></svg>\');background-size: 16px;}") as imsteve:
        gr.Label("Minecraft Skin Generator")
        gr.Markdown("Make AI generated Minecraft Skins by a Finetuned Stable Diffusion Version!<br>\nGithub Repository & Model used: https://github.com/Nick088Official/Minecraft_Skin_Generator<br>\n**Credits:**\n [Monadical-SAS](https://github.com/Monadical-SAS/minecraft_skin_generator)\n (Creators of the model), [Nick088](https://linktr.ee/Nick088) (Improving usage of the model)\n daroche (helping fix the 3d model texture isue)\n [Brottweiler](https://gist.github.com/Brottweiler/483d0856c6692ef70cf90bf1a85ce364)(script to fix the 3d model texture)\n [not-holar](https://huggingface.co/not-holar) (made the rendering of the image asset in the web ui look pixelated like minecraft and have a checkered background)\n[meew](https://huggingface.co/spaces/meeww/Minecraft_Skin_Generator/blob/main/models/player_model.glb) (Minecraft Player 3d model) <br>\n [![Discord](https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge)](https://discord.gg/AQsmBmgEPy)")

        with gr.Row():
            prompt = gr.Textbox(label="Your Prompt", info="What the Minecraft Skin should look like")
            stable_diffusion_model = gr.Dropdown(['2', 'xl'], value="xl", label="Stable Diffusion Model", info="Choose which Stable Diffusion Model to use, xl understands prompts better")
            model_precision_type = gr.Dropdown(["fp16", "fp32"], value="fp16", label="Model Precision Type", info="The precision type to load the model, like fp16 which is faster, or fp32 which is more precise but more resource consuming")

        with gr.Row():
            num_inference_steps = gr.Slider(label="Number of Inference Steps", info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference", minimum=1, maximum=50, value=25, step=1)
            guidance_scale = gr.Slider(label="Guidance Scale", info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.", minimum=0.0, maximum=10.0, value=7.5, step=0.1)
            seed = gr.Slider(value=42, minimum=0, maximum=MAX_SEED, step=1, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")
        
        filename = gr.Textbox(label="Output Image Name", info="The name of the file of the output image skin, keep the.png", value="output-skin.png")
        with gr.Row():
            model_3d = gr.Checkbox(label="See as 3D Model too", info="View the generated skin as a 3D Model too", value=True)
            verbose = gr.Checkbox(label="Verbose Output", info="Produce more detailed output while running", value=False)
        
        generate_skn = gr.Button("Generate")
        image_output = gr.Image(label="Generated Minecraft Skin Image Asset")
        image3d_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model View of the Skin")
        
        generate_skn.click(
            fn=run_inference,
            inputs=[
                prompt,
                stable_diffusion_model,
                num_inference_steps,
                guidance_scale,
                model_precision_type,
                seed,
                filename,
                model_3d,
                verbose,
            ],
            outputs=[
                image_output,
                image3d_output,
            ],
        )
    return imsteve

def main():
    # Handle Hugging Face Spaces environment
    if os.environ.get("HF_TOKEN") or os.environ.get("HF_HOME") or os.environ.get("HUGGINGFACE_HUB_CACHE"):
        import spaces
        @spaces.GPU(duration=75)
        def wrapped_create_gradio_ui():
            return create_gradio_ui()
        demo = wrapped_create_gradio_ui()
    else:
        demo = create_gradio_ui()

    setup_repository()
    demo.launch(show_api=True, share=True)

if __name__ == "__main__":
    main()