Initial Commit

This commit is contained in:
2026-05-14 10:23:47 +02:00
commit a655feff9e
8 changed files with 4357 additions and 0 deletions

View File

@@ -0,0 +1,262 @@
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 23 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()