"""Image generation application on modal.""" # # Stable Diffusion XL 1.0 # https://huggingface.co/docs/diffusers/main/en/using-diffusers/lcm # SDXL lightning # https://huggingface.co/ByteDance/SDXL-Lightning # Resources for NSFW # https://huggingface.co/spaces/radames/Real-Time-Text-to-Image-SDXL-Lightning/blob/main/app.py import datetime import json import os import tempfile import time import uuid from collections import deque from threading import Event, Lock, Thread import random from modal import App, Image, Retries, Secret, gpu, method from pathlib import Path from PIL import Image as PILImage from suno_utils.worker.schema import QueueItem from suno_utils.worker.tracing import distributed_trace, serialize_context from suno_utils.worker.utils import strip_square_brackets from suno_utils.worker.modal_base import MODAL_MOUNTS ############## CHANGE THESE ############## ENV_NAME = "dev" # dev, prod ########################################## APP_NAME = f"sdxl-{ENV_NAME}" # sdxl, stable-diffusion-xl-beta-4 assert APP_NAME.endswith(ENV_NAME) S3_BUCKET_NAME = "suno-data-uploads" S3_FOLDER_NAME = "studio/uploads" LARGE_IMAGE_PREFIX = "image_large_" IMAGE_PREFIX = "image_" IMAGE_WIDTH = 1024 IMAGE_HEIGHT = 1024 IMAGE_THUMBNAIL_WIDTH = 360 IMAGE_THUMBNAIL_HEIGHT = 360 # Set the number of concurrent requests to the model # GPU memory should be the limit # But threads slow it down... # Note that we are currently trading latency for throughput N_CONCURRENT = 2 # T4 - 1 ~ 10 GB, A10 - 2 ~ 15 GB N_CONCCURENT_BATCH = 2 N_CPU = 2 CONCURRENCY_LIMITS = { "dev": 5, "prod": 600, } KEEP_WARM = { "dev": 1, "prod": 60, } MOUNT_PATH = "/suno/models" random_seed = 42 if ENV_NAME == "dev" else int((time.time() * 1000) % 100000000) print("Random seed set to:", random_seed) random.seed(random_seed) ASSETS_PATH = Path(os.environ.get("SUNO_ASSETS_PATH", str((Path(__file__).parent / "assets")))) # ## Define a container image # # To take advantage of Modal's blazing fast cold-start times, we'll need to download our model weights # inside our container image with a download function. We ignore binaries, ONNX weights and 32-bit weights. # # Tip: avoid using global variables in this function to ensure the download step detects model changes and # triggers a rebuild. def download_models(): from huggingface_hub import hf_hub_download, snapshot_download ignore = ["*.bin", "*.onnx_data", "*/diffusion_pytorch_model.safetensors"] hf_download_path = hf_hub_download( "ByteDance/SDXL-Lightning", "sdxl_lightning_4step_unet.safetensors", local_dir=MOUNT_PATH ) print("downloaded to", hf_download_path) snapshot_download_path = snapshot_download( "stabilityai/stable-diffusion-xl-base-1.0", ignore_patterns=ignore ) print("downloaded to", snapshot_download_path) print("finish downloads") import transformers transformers.utils.move_cache() image = ( Image.debian_slim() .apt_install("libglib2.0-0", "libsm6", "libxrender1", "libxext6", "ffmpeg", "libgl1") .pip_install("openai>=1.12,<2") .pip_install("boto3") .pip_install( "diffusers~=0.29", "invisible_watermark~=0.1", "transformers>=4.31", "accelerate>=0.21", "safetensors>=0.3", "torch>=2.4", "ddtrace==2.8.0", "xformers>=0.0.22", ) .dockerfile_commands( [ "COPY --from=datadog/serverless-init /datadog-init /app/datadog-init", 'ENTRYPOINT ["/app/datadog-init"]', ] ) .run_function( download_models, secrets=[ Secret.from_name("huggingface-secret"), Secret.from_name("openai-secret"), ], ) .run_commands("WITH_CUDA=0 pip install stable-fast") ) app = App(APP_NAME, image=image) SECRETS = [ Secret.from_name("huggingface-secret"), Secret.from_name("openai-secret"), Secret.from_name("studio-aws"), Secret.from_name("api-callback-token"), Secret.from_dict( { "SUNO_ASSETS_PATH": "/suno/models/assets", } ), Secret.from_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": ENV_NAME, "DD_SERVICE": "sdxl-worker", "DD_LOGS_ENABLED": "true", "DD_TRACE_ENABLED": "true", }, ), Secret.from_name("datadog-metrics"), ] # just like engine, a thread in the background waiting for jobs # TODO: this could be offloaded to a CPU worker, for proper multiprocessing class StableDiffusionWorker(Thread): def __init__(self): start_time = time.time() # If reserved but unallocated memory is large try setting to avoid fragmentation. os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" print(os.environ["PYTORCH_CUDA_ALLOC_CONF"]) import torch from diffusers import EulerDiscreteScheduler, StableDiffusionXLPipeline, UNet2DConditionModel from safetensors.torch import load_file load_options = dict( torch_dtype=torch.float16, use_safetensors=True, variant="fp16", ) # Load base model unet = UNet2DConditionModel.from_config( "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet" ).to("cuda", torch.float16) # This apparently need a bit more polish... unet.load_state_dict( load_file( os.path.join(MOUNT_PATH, "sdxl_lightning_4step_unet.safetensors"), device="cuda", ) ) self.base = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", unet=unet, **load_options ) self.base.scheduler = EulerDiscreteScheduler.from_config( self.base.scheduler.config, timestep_spacing="trailing" ) self.base.to("cuda") print(f"start compiling, took {time.time() - start_time:.3f}s") from sfast.compilers.diffusion_pipeline_compiler import compile, CompilationConfig # sfast compilation compile_config = CompilationConfig.Default() import importlib.util xformers_spec = importlib.util.find_spec("xformers") if xformers_spec is not None: compile_config.enable_xformers = True triton_spec = importlib.util.find_spec("triton") if triton_spec is not None: compile_config.enable_triton = True # CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead. # But it can increase the amount of GPU memory used. # For StableVideoDiffusionPipeline it is not needed. compile_config.enable_cuda_graph = True self.base = compile(self.base, compile_config) print(f"finished compiling, took {time.time() - start_time:.3f}s") # sdxl engine setup self.add_request_lock = Lock() self.inference_batch_size = N_CONCCURENT_BATCH self.results = {} # placeholder for the output images, job_id -> image self.stop_flag = False self.doing_inference = False # flag to check if we are doing inference on GPU self.queue = deque() # queue of prompt, job_id self.max_queue_size = N_CONCURRENT # max size of the queue, same as the concurrent limit self.results_available = Event() import warnings warnings.filterwarnings("ignore", category=UserWarning) # start the thread as a daemon thread of current thread super().__init__(daemon=True) # warm up the model _ = self.inference("test") print(f"finished initializing, took {time.time() - start_time:.3f}s") def can_start_jobs(self): with self.add_request_lock: # has prompts, no results, not doing inference return len(self.queue) > 0 and len(self.results) == 0 and not self.doing_inference def add_prompt(self, prompt: str) -> int | None: """Add a prompt to the queue, return the job id. Job id is used to look up the result later. If the queue is full, return None. """ with self.add_request_lock: if len(self.queue) < self.max_queue_size: job_id = str(uuid.uuid4()) self.queue.append((prompt, job_id)) return job_id else: return None def get_prompt_result(self, job_index: int): while True: self.results_available.wait() with self.add_request_lock: if job_index not in self.results or self.doing_inference: # wait for the gpu job to finish pass else: # print("FOUND JOB DONE!") output_image = self.results.pop(job_index) if len(self.results) == 0: self.results_available.clear() return output_image # job is in queue but hasn't get the result # wait for a bit longer time.sleep(0.1) def run(self): while not self.stop_flag: if self.can_start_jobs(): # do a place hodler for the prompts with self.add_request_lock: self.doing_inference = True # clear the results again self.results = {} input_prompts = [] output_id_to_job_id = {} while self.queue and len(input_prompts) < self.inference_batch_size: prompt, job_index = self.queue.popleft() input_prompts.append(prompt) output_id_to_job_id[len(input_prompts) - 1] = job_index # now we just inference what we have on the prompts t0 = time.time() images = self.inference(input_prompts) total_time = time.time() - t0 print(f"Image gen of {len(input_prompts)} images took {total_time:.3f}s.") for output_id, image in enumerate(images): self.results[output_id_to_job_id[output_id]] = image self.doing_inference = False self.results_available.set() # print("finishing GPU inference, size", len(self.queue), len(self.results)) # queue is empty or results are not cleared, wait for a bit time.sleep(0.1) def inference(self, prompt: str | list[str], n_steps: int = 4) -> list: """Returns a list of generated images. For LCM, n_steps can be as low as 4. See https://huggingface.co/docs/diffusers/main/en/using-diffusers/lcm Nagative prompts don't have any effect in the denoising process. We will leave them here just for placebo effects. """ images = self.base( prompt=prompt, num_inference_steps=n_steps, guidance_scale=0, original_size=(IMAGE_WIDTH, IMAGE_HEIGHT), target_size=(IMAGE_WIDTH, IMAGE_HEIGHT), ).images return images # ## Load model and run inference # To avoid excessive cold-starts, we set the idle timeout to 240 seconds, meaning once a GPU has loaded the model it will stay # online for 4 minutes before spinning down. This can be adjusted for cost/experience trade-offs. @app.cls( # cpu=N_CPU, # dont require N_CPU to avoid backlog gpu=gpu.A10G(count=1), secrets=SECRETS, memory=15000, retries=Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), timeout=160, # compile can take a bit longer... container_idle_timeout=300, # 5 minutes, don't time out too short to avoid frequent worker spin up concurrency_limit=CONCURRENCY_LIMITS[ENV_NAME], allow_concurrent_inputs=N_CONCURRENT, keep_warm=KEEP_WARM[ENV_NAME], _experimental_buffer_containers=0 if ENV_NAME == "dev" else 12, region="us-east", ) class StableDiffusionInferencer: def __init__(self): import torch torch.set_num_threads(N_CPU) self.worker = StableDiffusionWorker() self.worker.start() # warm up the worker only for test run # _ = self.batched_inference("test") def batched_inference(self, prompt: str): # print("entered batch inference with prompt", prompt) # wait for the queue to be available while True: job_index = self.worker.add_prompt(prompt) if job_index is None: # queue is full --> sth is still running, waiting for a bit time.sleep(0.1) else: break # wait for inference -- should take ~ 2 seconds for 1 batch return self.worker.get_prompt_result(job_index) @method() @distributed_trace("batch_inference_and_upload_and_notify", "sdxl-worker", env_name=ENV_NAME) def batch_inference_and_upload_and_notify(self, prompt: str, item: QueueItem): curr_t0 = time.time() output_palette, output_concept = prompt.split(";") album_art_prompts = [ # 1st prompt f'Minimalist image showing "{output_concept}" ' f"using a retro film filter with grain " f"in an elevated and modern way " f'using "{output_palette}" colors.', # 2nd prompt f'Minimalist, whimsical image showing "{output_concept}" ' f"with a futuristic twist and a vintage film vibe, " f'with "{output_palette}" colors and grain textures.', ] new_prompt = random.choice(album_art_prompts) new_prompt += " No people or human elements visible in the frame." new_prompt += " Processed through a retro film filter with grain." print(f"ImageGen {item.id}: Using prompt: {prompt}") pil_image = self.batched_inference(new_prompt) total_time = time.time() - curr_t0 print(f"ImageGen {item.id}: Inference took {total_time:.3f}s ({(total_time):.3f}s / image).") return ImageUploader.upload_and_notify.spawn(pil_image, item, parent_context=serialize_context()) # This is a cpu worker that just uploads images to s3 @app.cls( cpu=1, secrets=SECRETS, timeout=100, container_idle_timeout=240, # 4 minutes, don't time out too short to avoid frequent worker spin up allow_concurrent_inputs=10, keep_warm=KEEP_WARM[ENV_NAME], mounts=MODAL_MOUNTS, cloud="aws", region="us-east", ) class ImageUploader: def _save_and_upload_image(self, pil_image, item: QueueItem) -> str: """Given a clip_id and an image, upload the image to s3, and return the s3 address.""" from suno_utils.worker.settings import s3_client clip_id = item.id # save things as jpegs with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: # overlay suno logo if item.is_bot_generation and random.random() < 0: # TADA logo = PILImage.open((ASSETS_PATH / "Logo-6.png").resolve()) logo_x = int(IMAGE_WIDTH * 0.3) logo_y = int(logo_x * 0.25) logo = logo.resize((logo_x, logo_y)) # center the logo. randomize it to make cropping of more difficult # if image is square we reshape it to 1/6 x -- 5/6 x random_x = random.randint(logo_x, int(IMAGE_WIDTH - logo_x - IMAGE_WIDTH * 0.2)) random_y = random.randint(logo_y, int(IMAGE_HEIGHT - logo_y * 1.25)) chosen_position = (random_x, random_y) # Hard-coded size, paste with transparency pil_image.paste(logo, chosen_position, mask=logo) fname = f"{LARGE_IMAGE_PREFIX}{clip_id}.jpeg" s3_fname = f"{S3_FOLDER_NAME}/{fname}" full_path = os.path.join(temp_dir, fname) pil_image.save(full_path, quality=75) # make a smaller image, for displays fname_small = f"{IMAGE_PREFIX}{clip_id}.jpeg" s3_fname_small = f"{S3_FOLDER_NAME}/{fname_small}" s3_small_path = f"s3://{S3_BUCKET_NAME}/{s3_fname_small}" pil_image_small = pil_image.resize((IMAGE_THUMBNAIL_WIDTH, IMAGE_THUMBNAIL_HEIGHT)) small_path = os.path.join(temp_dir, f"{clip_id}_small.jpeg") # will make the thumbnail even smaller pil_image_small.save(small_path, quality=50) s3_client.upload_file( full_path, S3_BUCKET_NAME, s3_fname, ExtraArgs={ "ContentType": "image/jpeg", }, ) s3_client.upload_file( small_path, S3_BUCKET_NAME, s3_fname_small, ExtraArgs={ "ContentType": "image/jpeg", }, ) print( "Uploaded to S3 to ", s3_small_path, ) return s3_small_path @method() @distributed_trace("upload_and_notify", "sdxl-worker", env_name=ENV_NAME) def upload_and_notify(self, pil_image, item: QueueItem) -> str: upload_link = self._save_and_upload_image(pil_image, item) item.notify_progress( { "id": item.id, "type": "image", "image_id": f"{IMAGE_PREFIX}{item.id}", }, ) return upload_link # This is a cpu worker that acts like a prompt moderator and image conductor @app.cls( cpu=1, secrets=SECRETS, timeout=100, container_idle_timeout=240, # 4 minutes, don't time out too short to avoid frequent worker spin up concurrency_limit=CONCURRENCY_LIMITS[ENV_NAME], allow_concurrent_inputs=10, _experimental_buffer_containers=0 if ENV_NAME == "dev" else 1, cloud="aws", region="us-east", ) class StableDiffusion: def generate_image_prompt(self, song_lyrics: str | None, do_sanitize: bool = False) -> str: import re from openai import OpenAI # since generate_image_prompt is not a modal method, you will need this key defined in local ENV # in order to test with `modal run` client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) song_lyrics = song_lyrics or "" song_lyrics = strip_square_brackets(song_lyrics) default_prompt = "3d soundwaves;dynamic, colorful" clean_image_generation_prompt = ( "You will receive some information about the song from the user. " "Your job is to use that information to summarize into two pieces of information: 'concept' and 'palette'. " "Clean the prompt by removing references to nudity, " "human body parts, sexual content, violence, hate, racism, or horror imagery. " "Make it PG-13 and ensure it won't produce scary or sexual images." "For 'concept', write a string that captures the cleaned image prompt. " "For 'palette' write a 1-3 word string about the colors that should be used, based on the feeling of the prompt. " """Output should be in the exact following json format: {"concept": "...", "palette": "..."} """ """DO NOT include the word `json` in the output.""" f"\n\n{song_lyrics}\n\n" ) lyrics_image_generation_prompt = ( "Your job is to help create a prompt for generating an album art picture for a song. " "You will receive some information about the song from the user. " "Your job is to use that information to summarize into two pieces of information: 'concept' and 'palette'. " "For 'concept', write a short string that helps depict something relevant to the song. " "Try to be specific like 'dog wearing sunglasses in the disco' for an edm song about someone's favorite dog. " "Avoid concepts that relate to Humans. " "Avoid proper nouns, people, human figures, nudity, sexual content, violence, or any harmful themes." "Make sure to capture enough information in the concept to not be ambiguous for downstream processing. " "For 'palette' write a 1-3 word string about the colors that should be used in a way that relates to the song and genre. " "Try to capture the mood of the song in the colors if possible, while sticking to a retro style. " "Input is in the following format: 'Lyrics: ... \nGenre: ...'" """Output should be in the exact following json format: {"concept": "...", "palette": "..."} """ """DO NOT include the word `json` in the output.""" f"\n\n{song_lyrics}\n\n" ) gpt_content_prompt = ( lyrics_image_generation_prompt if not do_sanitize else clean_image_generation_prompt ) # matching all languages given the prompt if not re.search(r"[\w]", song_lyrics): return default_prompt try: completion = client.chat.completions.create( model="gpt-4o-mini", messages=[ { "role": "user", "content": gpt_content_prompt, } ], max_tokens=800, ) output_prompt_info = completion.choices[0].message.content.split("\n\n")[-1] print(f"Output prompt info: {output_prompt_info}") output_prompt = json.loads(output_prompt_info) output_concept = output_prompt["concept"] output_palette = output_prompt["palette"] output_prompt = f"{output_palette};{output_concept}" print(f"Input prompt: {song_lyrics} \n Output prompt: {output_prompt}.") return output_prompt except Exception as e: print("Exception found!", e) return default_prompt @method() def generate_image(self, queue_item_json: str) -> str: """Single image generation API. For regenerating images with a text prompt. num_images and tags are kept for backwards compatability. """ item = QueueItem(**json.loads(queue_item_json)) prompt = self.generate_image_prompt(item.prompt_text, do_sanitize=True) print(f"ImageGen {item.id}: Single image generation with prompt: {prompt}") notify_call = StableDiffusionInferencer.batch_inference_and_upload_and_notify.spawn(prompt, item) image_url = notify_call.get(timeout=None) return image_url @method() @distributed_trace("generate_image_item", "sdxl-worker", env_name=ENV_NAME) def generate_image_item(self, queue_item_json: str) -> None: """Multi image generation API. Will generate multiple images based on the input queue item ids (clip_ids). """ item = QueueItem(**json.loads(queue_item_json)) print(f"ImageGen {item.ids}: Generating images for queue item {queue_item_json}.") t0 = time.time() if item.metadata.get("promotion") == "musical_memories": prompt = ( "An out of focus, low saturation, seventies style, day in the life, nostalgic, " f"scenic photograph for the memory {item.title or ''}, LUT, Kodak PORTRA 160 film, " "add grain effect" ) else: if item.metadata.get("gpt_description_prompt"): input_text_prompt = "Lyrics: " + (item.metadata.get("gpt_description_prompt") or "") else: input_text_prompt = ( "Lyrics: " + (item.prompt_text or "") + " \n Genre: " + (item.metadata.get("tags") or "instrumental") ) prompt = self.generate_image_prompt(input_text_prompt) total_time = time.time() - t0 print(f"ImageGen {item.ids}: ChatGPT took {total_time:.3f}s.") print(f"ImageGen {item.ids}: Multi image generation with prompt: {prompt}") if not item.ids: StableDiffusionInferencer.batch_inference_and_upload_and_notify.spawn( prompt, item, parent_context=serialize_context() ) else: # we need to make multiple image generation, with the same prompt! for clip_id in item.ids: new_item = item.copy(deep=True, update={"id": clip_id, "ids": None}) # note that we don't actually need the image url...they are implied as default _ = StableDiffusionInferencer.batch_inference_and_upload_and_notify.spawn( prompt, new_item, parent_context=serialize_context() ) time.sleep(0.001) # And this is our entrypoint; where the CLI is invoked. Explore CLI options # with: `modal run stable_diffusion_xl.py --prompt 'An astronaut riding a green horse'` def _test_generate_image_prompt(sd: StableDiffusion): print("testing generate_image_prompt") for _ in range(10): prompt = sd.generate_image_prompt( """I love watching sunsets with you It's my favorite thing in the world to do (ooh-yeah)""" ) print("prompt:", prompt) print("testing generate_image_prompt with null prompt") prompt = sd.generate_image_prompt(None) print("prompt:", prompt) @app.local_entrypoint() def main(prompt: str): sd = StableDiffusion() print(f"{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')} Done loading diffusion.") _test_generate_image_prompt(sd) t0 = time.time() num_generate_image_attempts = 24 for i in range(num_generate_image_attempts // 2): clip_ids = [str(i) + "_" + str(uuid.uuid4()), str(i) + "_" + str(uuid.uuid4())] new_queue_item = QueueItem(id="123", prompt_text=prompt, ids=clip_ids, metadata={}) sd.generate_image_item.spawn(new_queue_item.json()) time.sleep(0.01) total_time = time.time() - t0 print( f"{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')} \ submitting done --> Total took {total_time:.3f}s." ) time.sleep(200)