3d-model-generator/CLI/3d-model-image-prompt-generator.py

import os
import sys
import requests
import time
import json
import re
import subprocess

from langchain_community.llms import Ollama

SD_WEBUI_PATH = "/Users/giers/Tools/stable-diffusion-webui"

def is_sd_webui_running():
    try:
        r = requests.get("http://127.0.0.1:7860/sdapi/v1/txt2img", timeout=3)
        # Should error because POST is required, but if it responds, it's running
        return True
    except Exception:
        return False

def start_sd_webui_headless(webui_dir):
    # Use --headless, disable extensions/UIs for fastest startup
    args = [
        "python3", "launch.py",
        "--nowebui",         # Don't launch browser UI
        "--headless",        # No local UI window
        "--api",             # Enable API
        "--skip-torch-cuda-test",
        "--no-hashing",      # Faster startup
        "--disable-nan-check", # Optional: faster
        "--xformers"
    ]
    proc = subprocess.Popen(
        args,
        cwd=webui_dir,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    return proc

def get_output_dir():
    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
    os.makedirs(out_dir, exist_ok=True)
    return out_dir


def get_next_available_filename(base, ext=".png"):
    out_dir = get_output_dir()
    i = 1
    while True:
        fname = f"{base}-{i}{ext}"
        fpath = os.path.join(out_dir, fname)
        if not os.path.exists(fpath):
            return fpath
        i += 1

def flush_print(*args, **kwargs):
    print(*args, **kwargs)
    sys.stdout.flush()

OLLAMA_URL = "http://localhost:11434/api/generate"
#MODEL = "mistral-small3.1:24b"
MODEL = "mistral:latest"
SD_URL = "http://127.0.0.1:7860/sdapi/v1/txt2img"

META_PROMPT = """
You are an expert at writing concise, detailed Stable Diffusion prompts for 3D rendered objects as seen in professional game development. When I give you an object name, carefully follow these instructions step by step:

1. Subject ({{OBJECT}}):
- Start with the object name ({{OBJECT}}).
- Add 2–3 specific visual or material details (for example: shape, surface texture, design features like “chrome plating,” “organic armor,” “glowing eyes”).

2. Medium & Style:
- Add keywords for a 3D render and concept art look: 
  concept art, 3D render, game asset, professional game designer, digital sculpture, hyperrealistic, octane render, Unreal Engine 5, high poly.

3. Presentation:
- Explicitly state the object is shown fully in frame, not cropped, and completely visible from a neutral angle.
- No background: isolated on pure white background, or transparent background.

4. Lighting & Quality:
- Use studio lighting, no dramatic shadows, no depth of field, no blur.
- Emphasize sharp focus, ultra detailed, 8K, clean silhouette.

5. Artist Influence (optional):
- If fitting, add a well-known concept artist (example: by Beeple).

6. Negative Prompt:
- Add: Negative prompt: blurry, lowres, bad anatomy, distorted proportions, background, scenery, environment, artifacting, watermark, text, cropped, partial view, depth of field, out of frame, blur, vignette

Output Format:
- Line 1: All positive prompt keywords (subject, details, medium, style, lighting, quality, artist, presentation, background).
- Line 2: Start with Negative prompt: and then the negative keywords above.

Instructions to the LLM:
- The first line must always start with the object and its details.
- The object must be fully visible, not cropped, and entirely in the image frame.
- There must be no background, and the background must be pure white or transparent.
- Do NOT use blur, depth of field, vignette, or any visual effects that obscure details.
- Do NOT add any scenery or environment.
- Be concise and avoid repetition.

Now, generate a Stable Diffusion prompt for the object: {{OBJECT}}.
"""

SECOND_PROMPT_TEMPLATE = """
You are a prompt-to-JSON converter for image generation tasks.

Given a Stable Diffusion prompt, extract the positive and negative prompts.  
Also, analyze the subject:  
- If it is a single character, creature, full object, or person standing/upright, set "dimensions" to [512,768] (vertical).
- If it is a landscape, wide object, multi-character group, or scene, set "dimensions" to [768,512] (horizontal).
- If it is something square or best seen as a square (e.g., shield, logo, face, emblem, single centered item), set "dimensions" to [768,768] (square).

Return a JSON object with these fields:
{
  "positive prompt": "...",
  "negative prompt": "...",
  "dimensions": [W,H]
}

Respond only with the JSON.  
Here is the Stable Diffusion prompt:

{PROMPT}
"""


# --- streaming LLM utility ---
def stream_ollama(prompt):
    payload = {"model": MODEL, "prompt": prompt, "stream": True}
    with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=300) as r:
        r.raise_for_status()
        full_response = ""
        for line in r.iter_lines():
            if line:
                try:
                    chunk = json.loads(line.decode("utf-8"))["response"]
                except Exception:
                    chunk = line.decode("utf-8")
                print(chunk, end="", flush=True)
                full_response += chunk
        print()
    return full_response.strip()

def get_json_from_llm(prompt):
    ollama_chain = Ollama(model=MODEL, base_url="http://localhost:11434")
    response_gen = ollama_chain.stream(prompt)
    json_output = ""
    for chunk in response_gen:
        print(chunk, end="", flush=True)
        json_output += chunk
    print()
    # Try to extract JSON from output, even if LLM includes code block markers
    json_start = json_output.find("{")
    json_end = json_output.rfind("}")
    if json_start != -1 and json_end != -1:
        json_str = json_output[json_start:json_end+1]
        try:
            data = json.loads(json_str)
            return data
        except Exception as e:
            print("Error parsing JSON:", e)
            return None
    print("Failed to extract JSON!")
    return None

def sanitize_filename(s):
    # Replace spaces with underscores, remove non-alphanum chars
    s = s.strip().lower().replace(" ", "_")
    return re.sub(r'[^a-z0-9_]+', '', s)

def call_stable_diffusion(prompt_json, user_input):
    from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
    import torch

    # *** Use your diffusers-converted path! ***
    model_path = "/Volumes/SD/ML-Models/diffusers/dreamshaper_8_diffusers"

    if torch.backends.mps.is_available():
        device = "mps"
        print("Using Apple Silicon MPS backend")
    else:
        device = "cpu"
        print("Warning: Running on CPU (slow)")

    print("Loading model, this may take a while the first time...", flush=True)
    pipe = StableDiffusionPipeline.from_pretrained(
        model_path,
        torch_dtype=torch.float32,
        safety_checker=None,
        local_files_only=True,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to(device)
    pipe.set_progress_bar_config(disable=True)
    pipe.enable_attention_slicing()

    prompt = prompt_json["positive prompt"][:75]
    negative_prompt = prompt_json["negative prompt"][:75]
    width, height = prompt_json["dimensions"]
    num_images = 4
    steps = prompt_json.get("steps", 30)
    guidance_scale = prompt_json.get("cfg_scale", 6.5)

    print(f"Generating {num_images} image(s)...", flush=True)
    images = pipe(
        prompt=[prompt]*num_images,
        negative_prompt=[negative_prompt]*num_images,
        width=width,
        height=height,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
    ).images

    img_base = sanitize_filename(user_input)
    for idx, img in enumerate(images, start=1):
        filename = get_next_available_filename(f"{img_base}-{idx}")
        img.save(filename)
        print(f"Saved: {filename}")

def main():
    try:
        if len(sys.argv) < 2:
            print("Usage: python 3d-model-image-prompt-generator.py \"object name\"", flush=True)
            sys.exit(1)

        object_name = sys.argv[1]
        print(f"\n--- Generating Stable Diffusion prompt for: {object_name} ---\n", flush=True)
        prompt = META_PROMPT.replace("{{OBJECT}}", object_name)

        # 1. SD Prompt Generation (streamed)
        sd_prompt = stream_ollama(prompt)
        print("\n--- End of Stable Diffusion prompt ---\n", flush=True)

        # 2. JSON Conversion (streamed)
        print("--- Generating JSON for image generation ---\n", flush=True)
        second_prompt = SECOND_PROMPT_TEMPLATE.replace("{PROMPT}", sd_prompt)
        prompt_json = get_json_from_llm(second_prompt)
        if not prompt_json:
            print("Failed to get valid JSON. Exiting.", flush=True)
            sys.exit(1)

        # 3. Augment JSON with hard-coded SD settings
        prompt_json["checkpoint"] = "RealismPlus/dreamshaper_8.safetensors"
        prompt_json["vae"] = "Automatic"
        prompt_json["sampler"] = "DPM++ 2M Karras"
        prompt_json["steps"] = 30
        prompt_json["cfg_scale"] = 6.5

        print("\n--- Final prompt JSON ---\n", flush=True)
        print(json.dumps(prompt_json, indent=2, ensure_ascii=False), flush=True)

        # 4. Call local Stable Diffusion via diffusers
        call_stable_diffusion(prompt_json, object_name)
    except Exception as e:
        import traceback
        print("Exception in main():", e, flush=True)
        traceback.print_exc()

if __name__ == "__main__":
    main()