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)