File size: 1,314 Bytes
6026000 0f36a29 6026000 0f36a29 6026000 0f36a29 6026000 0f36a29 6026000 0f36a29 6026000 0f36a29 6026000 0f36a29 6026000 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import torch
from config import Config
class Generator:
def __init__(self, model_handler):
self.mh = model_handler
def predict(
self,
user_prompt,
negative_prompt="",
guidance_scale=1.2,
num_inference_steps=8,
seed=-1
):
# 1. Construct Prompt
if not user_prompt.strip():
# Fallback if user provides empty prompt
final_prompt = f"{Config.STYLE_TRIGGER}, a beautiful landscape, pixel art"
else:
final_prompt = f"{Config.STYLE_TRIGGER}, {user_prompt}"
print(f"Prompt: {final_prompt}")
# 2. Handle Seed
if seed == -1 or seed is None:
seed = torch.Generator().seed()
generator = torch.Generator(device=Config.DEVICE).manual_seed(int(seed))
print(f"Using seed: {seed}")
# 3. Run Text-to-Image Inference
print("Running pipeline...")
result = self.mh.pipeline(
prompt=final_prompt,
negative_prompt=negative_prompt,
generator=generator,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
clip_skip=2, # Optional, often helps with anime/pixel styles
).images[0]
return result |