init

app.py

@@ -1,26 +1,25 @@
-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 gradio as gr
+import spaces  # Necessário para o decorador @spaces.GPU (caso esteja usando Hugging Face Spaces)
 import os
+import torch
 from datetime import datetime
 from PIL import Image
 import boto3
 from botocore.exceptions import NoCredentialsError
 from dotenv import load_dotenv
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
 
 # Carregar variáveis de ambiente do arquivo .env
 load_dotenv()
 
-# AWS S3 Configuration
+# Configurações do AWS S3
 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
+HF_TOKEN = os.getenv('HF_TOKEN')  # Token da Hugging Face
 
-# Initialize S3 client
+# Inicializar cliente S3
 s3_client = boto3.client(
     's3',
     aws_access_key_id=AWS_ACCESS_KEY,
@@ -28,102 +27,100 @@ s3_client = boto3.client(
     region_name=AWS_REGION
 )
 
-# Configuration for the character pipeline
+# Configuração do pipeline para "character"
 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
+    use_auth_token=HF_TOKEN  # Inclui o token aqui
 )
 character_pipe.scheduler = EulerDiscreteScheduler.from_config(character_pipe.scheduler.config)
 
-# Configuration for the item pipeline
+# Configuração do pipeline para "item"
 item_pipe = DiffusionPipeline.from_pretrained(
     "openart-custom/DynaVisionXL",
     torch_dtype=torch.float16,
     use_safetensors=True,
-    use_auth_token=HF_TOKEN  # Include the token here
+    use_auth_token=HF_TOKEN  # Inclui o token aqui
 )
 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
+# Função de geração de imagem com alocação de GPU (através do decorador do Hugging Face Spaces)
+@spaces.GPU(duration=60)  # Aloca a GPU somente durante a execução desta função
 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]"
+        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'."
+        return "Tipo inválido. Escolha entre 'character' ou 'item'."
 
-    # Use custom prompts if provided
+    # Se o usuário fornecer prompt, utiliza-o; caso contrário, usa o padrão
     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
+    # Move o pipeline para a GPU
     pipe.to("cuda")
 
-    # Image generation
-    image = pipe(
+    # Geração da imagem
+    result = 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]
+    )
+    image = result.images[0]
 
-    # Save image to a temporary file
+    # Salva a imagem em um arquivo temporário
     temp_file = "/tmp/generated_image.png"
     image.save(temp_file)
 
-    # Upload to S3
+    # Faz upload para o AWS 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}
+        return "Credenciais não disponíveis"
+
+# Função que integra a geração via Gradio
+def gradio_generate(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
+    return generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps)
+
+# Definindo os componentes de entrada utilizando a API atual do Gradio
+model_type_input = gr.Dropdown(choices=["character", "item"], value="character", label="Model Type")
+prompt_input = gr.Textbox(lines=2, placeholder="Digite o prompt (deixe vazio para o padrão)", label="Prompt")
+negative_prompt_input = gr.Textbox(lines=2, placeholder="Digite o negative prompt (deixe vazio para o padrão)", label="Negative Prompt")
+width_input = gr.Number(value=512, label="Width")
+height_input = gr.Number(value=512, label="Height")
+guidance_scale_input = gr.Number(value=7.5, label="Guidance Scale")
+num_inference_steps_input = gr.Number(value=50, label="Number of Inference Steps")
+
+# Criação da interface Gradio
+iface = gr.Interface(
+    fn=gradio_generate,
+    inputs=[
+        model_type_input,
+        prompt_input,
+        negative_prompt_input,
+        width_input,
+        height_input,
+        guidance_scale_input,
+        num_inference_steps_input,
+    ],
+    outputs="text",
+    title="Image Generation API",
+    description="Gere imagens usando modelos de difusão e faça upload para o AWS S3."
+)
 
-# Run the FastAPI app with Uvicorn
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    iface.launch()

requirements.txt

@@ -1,10 +1,6 @@
-fastapi
-uvicorn
-transformers
-spaces
+gradio
 diffusers
 torch
 boto3
 python-dotenv
 Pillow
-accelerate