stablediffusion3_medium.py
1 import os 2 3 4 def get_python_executable(): 5 current_file_path = os.path.abspath(__file__) 6 project_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))) 7 8 return os.path.join(project_path, ".venvs/.venv-sd/bin/python") 9 10 def worker(prompt, sharedmem): 11 import base64 12 import io 13 from diffusers import StableDiffusion3Pipeline 14 import torch 15 from restai.config import RESTAI_DEFAULT_DEVICE 16 17 base = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16) 18 base.to(RESTAI_DEFAULT_DEVICE or "cuda") 19 20 image = base( 21 prompt=prompt, 22 negative_prompt="", 23 num_inference_steps=28, 24 guidance_scale=7.0, 25 ).images[0] 26 27 image_data = io.BytesIO() 28 image.save(image_data, format="JPEG") 29 image_base64 = base64.b64encode(image_data.getvalue()).decode('utf-8') 30 31 sharedmem["image"] = image_base64