Spaces:
Sleeping
Sleeping
| import os | |
| import requests | |
| from model_loading import GenerationSession | |
| def prompt_enhancer(user_prompt: str) -> str: | |
| try: | |
| response = requests.post( | |
| "http://localhost:11434/api/generate", | |
| json={ | |
| "model": "mistral", | |
| "format": "json", | |
| "prompt": ( | |
| f"[INST] You are an image generation prompt engineer. " | |
| f"Rewrite this prompt to be vivid and detailed, under 60 words. " | |
| f"Return ONLY the rewritten prompt, nothing else.\n\n" | |
| f"Prompt: {user_prompt} [/INST]" | |
| ), | |
| "stream": False | |
| }, | |
| timeout=60 | |
| ) | |
| except requests.exceptions.ConnectionError: | |
| print("Warning: Could not connect to local Ollama.") | |
| return user_prompt | |
| return response.json()["response"].strip() | |
| def smart_generate(user_prompt: str, session: GenerationSession, strength: float = 0.45): | |
| enhanced = prompt_enhancer(user_prompt) | |
| image = session.Generate(enhanced, strength=strength) | |
| return image[0], enhanced | |