from __future__ import annotations import io import os import re import time import tempfile from modal import Image, Secret, Stub, method from suno_utils.worker.modal_base import MODAL_MOUNTS from suno_utils.worker.utils import strip_square_brackets stub = Stub("stable-diffusion-cli") model_id = "runwayml/stable-diffusion-v1-5" cache_path = "/vol/cache" def download_models(): import diffusers import torch import openai hugging_face_token = os.environ["HUGGINGFACE_TOKEN"] openai.api_key = os.environ["OPENAI_API_KEY"] # Download scheduler configuration. Experiment with different schedulers # to identify one that works best for your use-case. scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained( model_id, subfolder="scheduler", use_auth_token=hugging_face_token, cache_dir=cache_path, ) scheduler.save_pretrained(cache_path, safe_serialization=True) # Downloads all other models. pipe = diffusers.StableDiffusionPipeline.from_pretrained( model_id, use_auth_token=hugging_face_token, revision="fp16", torch_dtype=torch.float16, cache_dir=cache_path, ) pipe.save_pretrained(cache_path, safe_serialization=True) image = ( Image.debian_slim(python_version="3.10") .pip_install("openai") .pip_install("boto3") .pip_install( "accelerate", "diffusers[torch]>=0.15.1", "ftfy", "torchvision", "transformers~=4.25.1", "triton", "safetensors", ) .pip_install( "torch==2.0.1+cu117", find_links="https://download.pytorch.org/whl/torch_stable.html", ) .pip_install("xformers", pre=True) .run_function( download_models, secrets=[ Secret.from_name("huggingface-secret"), Secret.from_name("openai-secret"), ], ) ) stub.image = image @stub.cls( gpu="A10G", secrets=[ Secret.from_name("huggingface-secret"), Secret.from_name("openai-secret"), Secret.from_name("studio-aws"), ], mounts=MODAL_MOUNTS, retries=3, ) class StableDiffusion: def __enter__(self): import diffusers import torch torch.backends.cuda.matmul.allow_tf32 = True scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained( cache_path, subfolder="scheduler", solver_order=2, prediction_type="epsilon", thresholding=False, algorithm_type="dpmsolver++", solver_type="midpoint", denoise_final=True, # important if steps are <= 10 low_cpu_mem_usage=True, device_map="auto", ) self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( cache_path, scheduler=scheduler, low_cpu_mem_usage=True, device_map="auto", ) self.pipe.enable_xformers_memory_efficient_attention() @method() def run_inference(self, prompt: str, steps: int = 10, batch_size: int = 5) -> list[bytes]: import torch with torch.inference_mode(): with torch.autocast("cuda"): images = self.pipe( [prompt] * batch_size, num_inference_steps=steps, guidance_scale=7.0, ).images # Convert to PNG bytes image_output = [] for image in images: with io.BytesIO() as buf: image.save(buf, format="PNG") image_output.append(buf.getvalue()) return image_output def generate_image_prompt(self, song_lyrics: str): import openai openai.api_key = os.environ["OPENAI_API_KEY"] song_lyrics = strip_square_brackets(song_lyrics) if not re.search(r"[a-zA-Z]", song_lyrics): return "Album art featuring a playful cartoon-style illustration of a cheerful dog" completion = openai.ChatCompletion.create( model="gpt-3.5-turbo", messages=[ { "role": "user", "content": f"What would simple cartoon style album art look like for the following song lyrics: '{song_lyrics}' Output at most 10 words of descriptive plain text. Your response should just include a single description of the image. Do not include any proper nouns or people.", } ], max_tokens=200, ) return completion.choices[0].message.content.split("\n\n")[-1] @method() def generate_image(self, song_lyrics: str, clip_id, num_images=1): from suno_utils.worker.settings import s3_client prompt = ( self.generate_image_prompt(song_lyrics) + " Output a simple cartoon style image with a blue background." ) print("Generating images with prompt: ", prompt) import uuid clip_id = str(uuid.uuid4()) with tempfile.TemporaryDirectory() as td: t0 = time.time() images = self.run_inference(prompt, batch_size=num_images) total_time = time.time() - t0 print(f"Sample took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).") for j, image_bytes in enumerate(images): fname = f"image_{clip_id}_{j}.png" full_path = os.path.join(td, fname) with open(full_path, "wb") as f: f.write(image_bytes) s3_client.upload_file(full_path, "suno-data-uploads", f"studio/uploads/{fname}") upload_link = f"s3://suno-data-uploads/studio/uploads/{fname}" print( "Uploaded to S3 to ", upload_link, ) return upload_link @stub.local_entrypoint() def main(): sd = StableDiffusion() song_lyrics = "[country banjo] Value (Error), Value (Error),\nValue (Error), Value (Error)\nLook at the server its in a blaze (whoa)\ntell all the coders to get away\nheard that they let the bugs out the cage (yeah)\nIf you sit and just listen you hear the page\nSend me the commit i'm on the case (case)\nTell em drop me the hash imma renovate\nI'm on Git Hub on my dinner date\nOn call, let me demonstrate\nWe need somebody to come back and save the day" prompt = sd.generate_image_prompt.call(song_lyrics) print("Generating images with prompt: ", prompt) import uuid clip_id = str(uuid.uuid4()) # sd.generate_image(song_lyrics, clip_id) # with tempfile.TemporaryDirectory() as td: # # if not dir.exists(): # # dir.mkdir(exist_ok=True, parents=True) # t0 = time.time() # images = sd.run_inference.call(prompt) # total_time = time.time() - t0 # print( # f"Sample took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)." # ) # for j, image_bytes in enumerate(images): # fname = f"image_{clip_id}_{j}.png" # full_path = os.path.join(td, fname) # # output_path = td / f"output_{j}.png" # print(f"Saving it to {full_path}") # with open(full_path, "wb") as f: # f.write(image_bytes) # s3_client.upload_file( # full_path, "suno-data-uploads", f"studio/uploads/{fname}.png" # ) # print("Uploaded to S3 to ", f"suno-data-uploads/studio/uploads/{fname}.png") sd.generate_image.call(song_lyrics, clip_id, num_images=1)