/ restai / image / workers / stablediffusion3_medium.py
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