File size: 4,680 Bytes
0505dd8 bf9983e | 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 | import gradio as gr
def greet(name):
return "Hello " + name + "!!"
demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import spaces # Necessary for the @spaces.GPU decorator
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
import torch
import os
from datetime import datetime
from PIL import Image
import boto3
from botocore.exceptions import NoCredentialsError
from dotenv import load_dotenv
# Carregar variáveis de ambiente do arquivo .env
load_dotenv()
# AWS S3 Configuration
AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')
AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
AWS_BUCKET_NAME = os.getenv('AWS_BUCKET_NAME')
AWS_REGION = os.getenv('AWS_REGION')
HF_TOKEN = os.getenv('HF_TOKEN') # Add this line to load your Hugging Face token
# Initialize S3 client
s3_client = boto3.client(
's3',
aws_access_key_id=AWS_ACCESS_KEY,
aws_secret_access_key=AWS_SECRET_KEY,
region_name=AWS_REGION
)
# Configuration for the character pipeline
character_pipe = DiffusionPipeline.from_pretrained(
"cagliostrolab/animagine-xl-3.1",
torch_dtype=torch.float16,
use_safetensors=True,
use_auth_token=HF_TOKEN # Include the token here
)
character_pipe.scheduler = EulerDiscreteScheduler.from_config(character_pipe.scheduler.config)
# Configuration for the item pipeline
item_pipe = DiffusionPipeline.from_pretrained(
"openart-custom/DynaVisionXL",
torch_dtype=torch.float16,
use_safetensors=True,
use_auth_token=HF_TOKEN # Include the token here
)
item_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(item_pipe.scheduler.config)
# Function for image generation with ZeroGPU
@spaces.GPU(duration=60) # Allocate GPU only during the execution of this function
def generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
if model_type == "character":
pipe = character_pipe
default_prompt = "1girl, souji okita, fate series, solo, upper body, bedroom, night, seducing, (sexy clothes)"
default_negative_prompt = "lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"
elif model_type == "item":
pipe = item_pipe
default_prompt = "great sword, runes on blade, acid on blade, weapon, (((item)))"
default_negative_prompt = "1girl, girl, man, boy, 1man, men, girls"
else:
return "Invalid type. Choose between 'character' or 'item'."
# Use custom prompts if provided
final_prompt = prompt if prompt else default_prompt
final_negative_prompt = negative_prompt if negative_prompt else default_negative_prompt
# Move the pipeline to the GPU
pipe.to("cuda")
# Image generation
image = pipe(
prompt=final_prompt,
negative_prompt=final_negative_prompt,
width=int(width),
height=int(height),
guidance_scale=float(guidance_scale),
num_inference_steps=int(num_inference_steps)
).images[0]
# Save image to a temporary file
temp_file = "/tmp/generated_image.png"
image.save(temp_file)
# Upload to S3
file_name = datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
try:
s3_client.upload_file(temp_file, AWS_BUCKET_NAME, file_name)
s3_url = f"https://{AWS_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{file_name}"
return s3_url
except NoCredentialsError:
return "Credentials not available"
# Initialize FastAPI
app = FastAPI()
# Define request model
class PredictRequest(BaseModel):
model_type: str
prompt: str = ""
negative_prompt: str = ""
width: int
height: int
guidance_scale: float
num_inference_steps: int
# Add FastAPI routes
@app.get("/")
def read_root():
return {"Hello World"}
@app.post("/api/predict")
async def predict(request: PredictRequest):
result = generate_image(
model_type=request.model_type,
prompt=request.prompt,
negative_prompt=request.negative_prompt,
width=request.width,
height=request.height,
guidance_scale=request.guidance_scale,
num_inference_steps=request.num_inference_steps
)
if result is None:
raise HTTPException(status_code=400, detail="Invalid input")
return {"result": result}
# Run the FastAPI app with Uvicorn
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |