262 lines
9.3 KiB
Python
262 lines
9.3 KiB
Python
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() |