Spaces:
Runtime error
Runtime error
| 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 | |
| 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) |