Spaces:
Runtime error
Runtime error
File size: 8,915 Bytes
431c7a8 8168982 431c7a8 8168982 431c7a8 8168982 431c7a8 8168982 431c7a8 8168982 431c7a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 |
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional
import logging
import os
import boto3
import json
import shlex
import subprocess
import tempfile
import time
import base64
import gradio as gr
import numpy as np
import rembg
import spaces
import torch
from PIL import Image
from functools import partial
import io
from io import BytesIO
app = FastAPI()
torch.cuda.empty_cache()
subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
HEADER = """FRAME AI"""
if torch.cuda.is_available():
device = "cuda:0"
else:
device = "cpu"
model = TSR.from_pretrained(
"stabilityai/TripoSR",
config_name="config.yaml",
weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)
rembg_session = rembg.new_session()
ACCESS = os.getenv("ACCESS")
SECRET = os.getenv("SECRET")
bedrock = boto3.client(service_name='bedrock', aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
bedrock_runtime = boto3.client(service_name='bedrock-runtime', aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
# def generate_image_from_text(pos_prompt):
# # bedrock_runtime = boto3.client(region_name = 'us-east-1', service_name='bedrock-runtime')
# parameters = {'text_prompts': [{'text': pos_prompt , 'weight':1},
# {'text': """Blurry, out of frame, out of focus, Detailed, dull, duplicate, bad quality, low resolution, cropped""", 'weight': -1}],
# 'cfg_scale': 7, 'seed': 0, 'samples': 1}
# request_body = json.dumps(parameters)
# response = bedrock_runtime.invoke_model(body=request_body,modelId = 'stability.stable-diffusion-xl-v1')
# response_body = json.loads(response.get('body').read())
# base64_image_data = base64.b64decode(response_body['artifacts'][0]['base64'])
# return Image.open(io.BytesIO(base64_image_data))
def gen_pos_prompt(text):
instruction = f'''Your task is to create a positive prompt for image generation.
Objective: Generate images that prioritize structural integrity and accurate shapes. The focus should be on the correct form and basic contours of objects, with minimal concern for colors.
Guidelines:
Complex Objects (e.g., animals, vehicles): For these, the image should resemble a toy object, emphasizing the correct shape and structure while minimizing details and color complexity.
Example Input: A sports bike
Example Positive Prompt: Simple sports bike with accurate shape and structure, minimal details, digital painting, concept art style, basic contours, soft lighting, clean lines, neutral or muted colors, toy-like appearance, low contrast.
Example Input: A lion
Example Positive Prompt: Toy-like depiction of a lion with a focus on structural accuracy, minimal details, digital painting, concept art style, basic contours, soft lighting, clean lines, neutral or muted colors, simplified features, low contrast.
Simple Objects (e.g., a tennis ball): For these, the prompt should specify a realistic depiction, focusing on the accurate shape and structure.
Example Input: A tennis ball
Example Positive Prompt: Realistic depiction of a tennis ball with accurate shape and texture, digital painting, clean lines, minimal additional details, soft lighting, neutral or muted colors, focus on structural integrity.
Prompt Structure:
Subject: Clearly describe the object and its essential shape and structure.
Medium: Specify the art style (e.g., digital painting, concept art).
Style: Include relevant style terms (e.g., simplified, toy-like for complex objects; realistic for simple objects).
Resolution: Mention resolution if necessary (e.g., basic resolution).
Lighting: Indicate the type of lighting (e.g., soft lighting).
Color: Use neutral or muted colors with minimal emphasis on color details.
Additional Details: Keep additional details minimal or specify if not desired.
Input: {text}
Positive Prompt:
'''
body = json.dumps({'inputText': instruction,
'textGenerationConfig': {'temperature': 0.1, 'topP': 0.01, 'maxTokenCount':512}})
response = bedrock_runtime.invoke_model(body=body, modelId='amazon.titan-text-express-v1')
pos_prompt = json.loads(response.get('body').read())['results'][0]['outputText']
return pos_prompt
def encode_image_to_base64(image):
with io.BytesIO() as buffered:
image.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode('utf-8')
def generate_image_from_text(encoded_image, seed, pos_prompt=None):
neg_prompt = '''Detailed, complex textures, intricate patterns, realistic lighting, high contrast, reflections, fuzzy surface, realistic proportions, photographic quality, vibrant colors, detailed background, shadows, disfigured, deformed, ugly, multiple, duplicate.'''
encoded_str = encode_image_to_base64(encoded_image)
if pos_prompt:
parameters = {
'taskType': 'IMAGE_VARIATION',
'imageVariationParams': {
'images': [encoded_str],
'text': gen_pos_prompt(pos_prompt),
'negativeText': neg_prompt,
'similarityStrength': 0.7
},
'imageGenerationConfig': {
"cfgScale": 8,
"seed": int(seed),
"width": 512,
"height": 512,
"numberOfImages": 1
}
}
else:
parameters = {
'taskType': 'IMAGE_VARIATION',
'imageVariationParams': {
'images': [encoded_str],
'negativeText': neg_prompt,
'similarityStrength': 0.7
},
'imageGenerationConfig': {
"cfgScale": 8,
"seed": int(seed),
"width": 512,
"height": 512,
"numberOfImages": 1
}
}
request_body = json.dumps(parameters)
response = bedrock_runtime.invoke_model(body=request_body, modelId='amazon.titan-image-generator-v1')
response_body = json.loads(response.get('body').read())
base64_image_data = base64.b64decode(response_body['images'][0])
return Image.open(io.BytesIO(base64_image_data))
def check_input_image(input_image):
if input_image is None:
raise gr.Error("No image uploaded!")
def preprocess(input_image, do_remove_background, foreground_ratio):
def fill_background(image):
image = np.array(image).astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = Image.fromarray((image * 255.0).astype(np.uint8))
return image
if do_remove_background:
image = input_image.convert("RGB")
image = remove_background(image, rembg_session)
image = resize_foreground(image, foreground_ratio)
image = fill_background(image)
else:
image = input_image
if image.mode == "RGBA":
image = fill_background(image)
return image
# @spaces.GPU
def generate(image, mc_resolution, formats=["obj", "glb"]):
scene_codes = model(image, device=device)
mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
mesh = to_gradio_3d_orientation(mesh)
mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
mesh.export(mesh_path_glb.name)
mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
mesh.apply_scale([-1, 1, 1]) # Otherwise the visualized .obj will be flipped
mesh.export(mesh_path_obj.name)
return mesh_path_obj.name, mesh_path_glb.name
@app.post("/process_image/")
async def process_image(
file: UploadFile = File(...),
seed: int = Form(...),
use_image: bool = Form(...),
do_remove_background: bool = Form(...),
foreground_ratio: float = Form(...),
mc_resolution: int = Form(...),
text_prompt: Optional[str] = Form(None)
):
image_bytes = await file.read()
input_image = Image.open(BytesIO(image_bytes))
if use_image:
image_pil = generate_image_from_text(encoded_image=input_image, seed=seed, pos_prompt=text_prompt)
else:
image_pil = input_image
preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
mesh_name_obj, mesh_name_glb = generate(preprocessed, mc_resolution)
return {
"obj_path": mesh_name_obj,
"glb_path": mesh_name_glb
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |