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()