"""Image generation application on modal. Uses flux. More expensive than sdxl lightning. (A100 vs A10G) Careful deploying it. """ # # Stable Diffusion XL 1.0 # https://huggingface.co/docs/diffusers/main/en/using-diffusers/lcm # Flux # https://huggingface.co/black-forest-labs/FLUX.1-schnell # https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux # Resources for NSFW # https://huggingface.co/spaces/radames/Real-Time-Text-to-Image-SDXL-Lightning/blob/main/app.py # Even modal # https://modal.com/docs/examples/flux import datetime import json import os import tempfile import time import uuid from collections import deque from threading import Event, Lock, Thread import random import modal from modal import App, Image, Queue, Retries, Secret, method from pathlib import Path from PIL import Image as PILImage import torch 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_model_volume import model_store_volume ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## APP_NAME = f"flux-{DEPLOYMENT_TYPE}" assert APP_NAME.endswith(DEPLOYMENT_TYPE) 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 # TODO: increase to at least 2 once modal scaling is fixed N_CONCURRENT = 8 # T4 - 1 ~ 10 GB, A10 - 2 ~ 15 GB N_CONCCURENT_BATCH = 4 # tweaks N_CPU = 4 CONCURRENCY_LIMITS = { "dev": 5, "prod": 100, # cap prod flux at 100 } KEEP_WARM = { "dev": 1, "prod": 15, # need to keep the image high cause warm up is long } MOUNT_PATH = "/suno/models" random_seed = 42 if DEPLOYMENT_TYPE == "dev" else int((time.time() * 1000) % 100000000) print("Random seed set to:", random_seed) random.seed(random_seed) # because this file is not in a sub directory in modal, we need to set the full path to the assets ASSETS_PATH = Path(str((Path(__file__).parent / "suno_utils/worker/assets"))) EVENTS_QUEUE_NAME = f"events-queue-{DEPLOYMENT_TYPE}" # ## 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(): import torch from diffusers import FluxPipeline local_model_path = os.path.join(MOUNT_PATH, "FLUX.1-schnell") if os.path.exists(local_model_path): pipe = FluxPipeline.from_pretrained( local_model_path, torch_dtype=torch.bfloat16, local_files_only=True ) else: pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16 ) pipe.save_pretrained(local_model_path) print(f"Download finished! {os.listdir(MOUNT_PATH)}") 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"]', ] ) .pip_install("diffusers~=0.30.2") .pip_install("transformers", "accelerate", "sentencepiece") .pip_install("ddtrace==2.21.1") .run_commands("WITH_CUDA=0 pip install stable-fast") .add_local_python_source("suno_utils", copy=False) ) app = App(APP_NAME, image=image) SECRETS = [ Secret.from_name("huggingface-secret-suno"), Secret.from_name("openai-secret"), Secret.from_name("studio-aws"), Secret.from_name("api-callback-token"), Secret.from_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": DEPLOYMENT_TYPE, "DD_SERVICE": "flux-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 FluxPipeline local_model_path = os.path.join(MOUNT_PATH, "FLUX.1-schnell") self.base = FluxPipeline.from_pretrained( local_model_path, torch_dtype=torch.bfloat16, local_files_only=True ) self.base.to("cuda") # 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) def optimize(self, pipe, compile=True): if not compile: return pipe # tag the compute-intensive modules, the Transformer and VAE decoder, for compilation pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True) pipe.vae.decode = torch.compile(pipe.vae.decode, mode="reduce-overhead", fullgraph=True) return pipe 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 warmup(self): job_ids = [] with self.add_request_lock: for _ in range(self.inference_batch_size + 1): job_id = str(uuid.uuid4()) self.queue.append(("test", job_id)) job_ids.append(job_id) for job_id in job_ids: result = self.get_prompt_result(job_id) print(f"Warmup job {job_id} finished with result {result}") 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): self.base = self.optimize(self.base, compile=True) 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 = {} # wait for a bit to collect more prompts # this shouldnt be necessary, but it signals to modal that 1 concurrency is inefficient # batch_start_time = time.time() while self.queue and len(input_prompts) < self.inference_batch_size: if self.queue: prompt, job_index = self.queue.popleft() input_prompts.append(prompt) output_id_to_job_id[len(input_prompts) - 1] = job_index else: time.sleep(0.1) # pad to the batch size while len(input_prompts) < self.inference_batch_size: input_prompts.append("bird") output_id_to_job_id[len(input_prompts) - 1] = -1 # 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): if output_id_to_job_id[output_id] != -1: 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.""" prompt = [prompt] if isinstance(prompt, str) else prompt images = self.base( prompt=prompt, num_inference_steps=n_steps, guidance_scale=0.0, height=IMAGE_HEIGHT, # needs to be divisable by 8 width=IMAGE_WIDTH, # needs to be divisable by 8 max_sequence_length=256, ).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, gpu=["H100"], # GPU fallbacks secrets=SECRETS, memory=15000, retries=Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), timeout=360, # compile can take a bit longer... scaledown_window=300, # 5 minutes, don't time out too short to avoid frequent worker spin up max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=KEEP_WARM[DEPLOYMENT_TYPE], buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 3, volumes={ "/suno": model_store_volume, }, # region="us-east", ) @modal.concurrent(max_inputs=N_CONCURRENT) class StableDiffusionInferencer: def __init__(self): import torch torch.set_num_threads(N_CPU) self.worker = StableDiffusionWorker() self.worker.start() print("Warming up the worker") self.worker.warmup() print("Worker warmed up") 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.05) 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", "flux-worker", env_name=DEPLOYMENT_TYPE) 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." 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=240, scaledown_window=240, # 4 minutes, don't time out too short to avoid frequent worker spin up min_containers=100 if DEPLOYMENT_TYPE == "prod" else 4, max_containers=1000, # cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=4) class ImageUploader: def __init__(self): self.events_queue = Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True) 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 with PILImage.open((ASSETS_PATH / "Logo-6.png").resolve()) as logo: 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", }, ) pil_image_small.close() pil_image.close() print( "Uploaded to S3 to ", s3_small_path, ) return s3_small_path @method() @distributed_trace("upload_and_notify", "flux-worker", env_name=DEPLOYMENT_TYPE) 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}", }, ) self.events_queue.put( { "type": "image_generated", "data": {"image_id": f"{IMAGE_PREFIX}{item.id}"}, }, partition=item.id, partition_ttl=60, block=False, ) 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, scaledown_window=240, # 4 minutes, don't time out too short to avoid frequent worker spin up max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=10) class StableDiffusion: def __init__(self): self.modal_inference_and_upload_and_notify = modal.Cls.from_name( f"flux-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "StableDiffusionInferencer", )().batch_inference_and_upload_and_notify @method() 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 def _create_prompt_from_item(self, item, sanitize=False): """Helper to create a consistent prompt from available item data.""" # For musical memories promotion, use the special prompt if item.metadata.get("promotion") == "musical_memories": return ( "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" ) # Create a single template that uses the available information template = f"""Lyrics: {item.prompt_text or ""} Title: {item.title or ""} Genre: {item.metadata.get("tags", "")}""" return self.generate_image_prompt.local(template, do_sanitize=sanitize) @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._create_prompt_from_item(item, sanitize=True) print(f"ImageGen {item.id}: Single image generation with prompt: {prompt}") notify_call = self.modal_inference_and_upload_and_notify.spawn(prompt, item) image_url = notify_call.get(timeout=None) return image_url @method() @distributed_trace("generate_image_item", "flux-worker", env_name=DEPLOYMENT_TYPE) 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() prompt = self._create_prompt_from_item(item) total_time = time.time() - t0 print(f"ImageGen {item.ids}: Prompt generation took {total_time:.3f}s.") print(f"ImageGen {item.ids}: Multi image generation with prompt: {prompt}") if not item.ids: self.modal_inference_and_upload_and_notify.spawn(prompt, item) 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 _ = self.modal_inference_and_upload_and_notify.spawn(prompt, new_item) 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.remote( """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.remote(None) print("prompt:", prompt) def _stress_test_image_generation(sd: StableDiffusion, prompt: str): print("Starting stress test - generating 1000 images") num_images = 1000 batch_size = 2 # Generate 2 images per queue item num_batches = num_images // batch_size for i in range(num_batches): clip_ids = [str(i) + "_" + str(uuid.uuid4()) for _ in range(batch_size)] new_queue_item = QueueItem(id="stress_test", prompt_text=prompt, ids=clip_ids, metadata={}) sd.generate_image_item.spawn(new_queue_item.json()) time.sleep(1e6) @app.local_entrypoint() def main(prompt: str = "A beautiful sunset over a calm ocean"): sd = StableDiffusion() print(f"{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')} Done loading diffusion.") _stress_test_image_generation(sd, prompt) _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(480)