diff --git a/modal_code.py b/modal_code.py index 76cf7f5..4c6c7b6 100644 --- a/modal_code.py +++ b/modal_code.py @@ -1,10 +1,13 @@ -import io +import os, io from pathlib import Path import random import uuid from fastapi import FastAPI, Request +from oxen import RemoteRepo +import oxen +import modal -from modal import Image, Mount, Stub, asgi_app, build, enter, gpu, method +from modal import Image, Secret, Stub, asgi_app, build, enter, gpu, method sdxl_image = ( Image.debian_slim() @@ -17,6 +20,7 @@ "transformers~=4.31", "accelerate~=0.21", "safetensors~=0.3", + "oxenai", ) ) @@ -30,7 +34,11 @@ from fastapi.responses import Response from huggingface_hub import snapshot_download +@stub.function(image=sdxl_image) def generate_random_prompt(username): + ''' + generate the + ''' adjectives = ["grand", "noble", "splendid", "majestic", "regal", "royal", "stately", "august", "magnificent", "dignified", "imposing", "sublime", "resplendent", "glorious", "impressive", "exalted", @@ -61,10 +69,18 @@ def generate_random_prompt(username): print(prompt) return prompt - +@stub.function(image=sdxl_image, secret=modal.Secret.from_name("oxenai-api-key")) def save_to_oxen(repo_name, branch_name, prompt, file_names, commit_message): + ''' + save the file to oxen repo + ''' + repo = RemoteRepo(repo_name) repo.checkout(branch_name) + + api_key = os.environ["OXENAI_API_KEY"] + oxen.auth.config_auth(api_key) + oxen.user.config_user("Bessie", "ox@oxen.ai") # Add all files in one commmit paths = [] @@ -128,14 +144,16 @@ def imagine(self, prompt, num_images=1, num_steps=50, width=1024, height=1024): num_inference_steps=num_steps ).images # type: ignore print(images) - file_names = [] + image_byte_list = [] + os.makedirs("images", exist_ok=True) for image in images: - file_name = "-".join(file_base_name) + "-" + str(uuid.uuid4()) + ".png" - image.save(file_name) - file_names.append(file_name) - # print(file_names) - return file_names + byte_stream = io.BytesIO() + image.save(byte_stream, format="PNG") + image_bytes = byte_stream.getvalue() + image_byte_list.append(image_bytes) + + return image_byte_list app = FastAPI() @@ -147,19 +165,55 @@ async def generate_image( data = await request.json() print(f"POST /imagine - received data={data}") - if data['action'] == 'created': - data['prompt'] = generate_random_prompt(data['sender']['login']) + if 'action' in data and data['action'] == 'created': + data['prompt'] = generate_random_prompt.remote(data['sender']['login']) if data['prompt']: prompt = data['prompt'] - file_names = Model().imagine.remote(prompt) + image_byte_list = Model().imagine.remote(prompt) print("Saving to Oxen") commit_message = f"{prompt}" namespace = "ox" - repo_name = "GenerativeImagePlayground" + repo_name = "FlyingOxen" branch_name = "main" - commit, images = save_to_oxen(f"{namespace}/{repo_name}", branch_name, prompt, file_names, commit_message) + + repo = RemoteRepo(f"{namespace}/{repo_name}") + repo.checkout(branch_name) + + api_key = os.environ["OXENAI_API_KEY"] + oxen.auth.config_auth(api_key) + oxen.user.config_user("Bessie", "ox@oxen.ai") + + # Add all files in one commmit + paths = [] + os.makedirs("images", exist_ok=True) + for image_bytes in image_byte_list: + file_base_name = str(prompt).split(" ")[:10] + basename = "-".join(file_base_name) + "-" + str(uuid.uuid4()) + ".png" + file_name = os.path.join("images", basename) + + print(f"Saving it to {file_name}") + with open(file_name, "wb") as f: + f.write(image_bytes) + + repo.add(file_name, "images") + repo_path = f"images/{file_name}" + paths.append(repo_path) + repo.add_df_row("annotations.jsonl", {"prompt": prompt, "image": repo_path, "model": "stable-diffusion-xl-base-1.0"}) + + commit = repo.commit(commit_message) + images = [] + for path in paths: + images.append({ + "path": path, + "content_url": f"https://hub.oxen.ai/api/repos/{namespace}/{repo_name}/file/{commit.commit_id}/{path}", + # https://www.oxen.ai/ox/GenerativeImagePlayground/file/edfdd2bc11396d86/images/a-white-fluffy-ox-427b3fa6-fed2-4379-b9a8-04babf476bc8.png + "view_url": f"https://oxen.ai/{namespace}/{repo_name}/file/{commit.commit_id}/{path}", + }) + print(f"Got Commit: {commit}") + os.remove(file_name) + commit_id = commit.commit_id return { "status": "success", @@ -174,13 +228,12 @@ async def generate_image( else: return {"status": "error", "status_message": "Must supply 'prompt' field in json"} - @app.post("/echo") async def foo(request: Request): data = await request.json() return data -@stub.function(image=sdxl_image) +@stub.function(image=sdxl_image, secret=modal.Secret.from_name("oxenai-api-key")) @asgi_app() def fastapi_app(): return app