Devsora commited on
Commit
c7888bc
·
1 Parent(s): 82376eb
Files changed (2) hide show
  1. app.py +57 -60
  2. requirements.txt +1 -5
app.py CHANGED
@@ -1,26 +1,25 @@
1
- from fastapi import FastAPI, HTTPException
2
- from pydantic import BaseModel
3
- import spaces # Necessary for the @spaces.GPU decorator
4
- from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
5
- import torch
6
  import os
 
7
  from datetime import datetime
8
  from PIL import Image
9
  import boto3
10
  from botocore.exceptions import NoCredentialsError
11
  from dotenv import load_dotenv
 
12
 
13
  # Carregar variáveis de ambiente do arquivo .env
14
  load_dotenv()
15
 
16
- # AWS S3 Configuration
17
  AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')
18
  AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
19
  AWS_BUCKET_NAME = os.getenv('AWS_BUCKET_NAME')
20
  AWS_REGION = os.getenv('AWS_REGION')
21
- HF_TOKEN = os.getenv('HF_TOKEN') # Add this line to load your Hugging Face token
22
 
23
- # Initialize S3 client
24
  s3_client = boto3.client(
25
  's3',
26
  aws_access_key_id=AWS_ACCESS_KEY,
@@ -28,102 +27,100 @@ s3_client = boto3.client(
28
  region_name=AWS_REGION
29
  )
30
 
31
- # Configuration for the character pipeline
32
  character_pipe = DiffusionPipeline.from_pretrained(
33
  "cagliostrolab/animagine-xl-3.1",
34
  torch_dtype=torch.float16,
35
  use_safetensors=True,
36
- use_auth_token=HF_TOKEN # Include the token here
37
  )
38
  character_pipe.scheduler = EulerDiscreteScheduler.from_config(character_pipe.scheduler.config)
39
 
40
- # Configuration for the item pipeline
41
  item_pipe = DiffusionPipeline.from_pretrained(
42
  "openart-custom/DynaVisionXL",
43
  torch_dtype=torch.float16,
44
  use_safetensors=True,
45
- use_auth_token=HF_TOKEN # Include the token here
46
  )
47
  item_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(item_pipe.scheduler.config)
48
 
49
- # Function for image generation with ZeroGPU
50
- @spaces.GPU(duration=60) # Allocate GPU only during the execution of this function
51
  def generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
52
  if model_type == "character":
53
  pipe = character_pipe
54
  default_prompt = "1girl, souji okita, fate series, solo, upper body, bedroom, night, seducing, (sexy clothes)"
55
- 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]"
 
 
56
  elif model_type == "item":
57
  pipe = item_pipe
58
  default_prompt = "great sword, runes on blade, acid on blade, weapon, (((item)))"
59
  default_negative_prompt = "1girl, girl, man, boy, 1man, men, girls"
60
  else:
61
- return "Invalid type. Choose between 'character' or 'item'."
62
 
63
- # Use custom prompts if provided
64
  final_prompt = prompt if prompt else default_prompt
65
  final_negative_prompt = negative_prompt if negative_prompt else default_negative_prompt
66
 
67
- # Move the pipeline to the GPU
68
  pipe.to("cuda")
69
 
70
- # Image generation
71
- image = pipe(
72
  prompt=final_prompt,
73
  negative_prompt=final_negative_prompt,
74
  width=int(width),
75
  height=int(height),
76
  guidance_scale=float(guidance_scale),
77
  num_inference_steps=int(num_inference_steps)
78
- ).images[0]
 
79
 
80
- # Save image to a temporary file
81
  temp_file = "/tmp/generated_image.png"
82
  image.save(temp_file)
83
 
84
- # Upload to S3
85
  file_name = datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
86
  try:
87
  s3_client.upload_file(temp_file, AWS_BUCKET_NAME, file_name)
88
  s3_url = f"https://{AWS_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{file_name}"
89
  return s3_url
90
  except NoCredentialsError:
91
- return "Credentials not available"
92
-
93
- # Initialize FastAPI
94
- app = FastAPI()
95
-
96
- # Define request model
97
- class PredictRequest(BaseModel):
98
- model_type: str
99
- prompt: str = ""
100
- negative_prompt: str = ""
101
- width: int
102
- height: int
103
- guidance_scale: float
104
- num_inference_steps: int
105
-
106
- # Add FastAPI routes
107
- @app.get("/")
108
- def read_root():
109
- return {"Hello": "World"}
110
-
111
- @app.post("/api/predict")
112
- async def predict(request: PredictRequest):
113
- result = generate_image(
114
- model_type=request.model_type,
115
- prompt=request.prompt,
116
- negative_prompt=request.negative_prompt,
117
- width=request.width,
118
- height=request.height,
119
- guidance_scale=request.guidance_scale,
120
- num_inference_steps=request.num_inference_steps
121
- )
122
- if result is None:
123
- raise HTTPException(status_code=400, detail="Invalid input")
124
- return {"result": result}
125
 
126
- # Run the FastAPI app with Uvicorn
127
  if __name__ == "__main__":
128
- import uvicorn
129
- uvicorn.run(app, host="0.0.0.0", port=7860)
 
1
+ import gradio as gr
2
+ import spaces # Necessário para o decorador @spaces.GPU (caso esteja usando Hugging Face Spaces)
 
 
 
3
  import os
4
+ import torch
5
  from datetime import datetime
6
  from PIL import Image
7
  import boto3
8
  from botocore.exceptions import NoCredentialsError
9
  from dotenv import load_dotenv
10
+ from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
11
 
12
  # Carregar variáveis de ambiente do arquivo .env
13
  load_dotenv()
14
 
15
+ # Configurações do AWS S3
16
  AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')
17
  AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
18
  AWS_BUCKET_NAME = os.getenv('AWS_BUCKET_NAME')
19
  AWS_REGION = os.getenv('AWS_REGION')
20
+ HF_TOKEN = os.getenv('HF_TOKEN') # Token da Hugging Face
21
 
22
+ # Inicializar cliente S3
23
  s3_client = boto3.client(
24
  's3',
25
  aws_access_key_id=AWS_ACCESS_KEY,
 
27
  region_name=AWS_REGION
28
  )
29
 
30
+ # Configuração do pipeline para "character"
31
  character_pipe = DiffusionPipeline.from_pretrained(
32
  "cagliostrolab/animagine-xl-3.1",
33
  torch_dtype=torch.float16,
34
  use_safetensors=True,
35
+ use_auth_token=HF_TOKEN # Inclui o token aqui
36
  )
37
  character_pipe.scheduler = EulerDiscreteScheduler.from_config(character_pipe.scheduler.config)
38
 
39
+ # Configuração do pipeline para "item"
40
  item_pipe = DiffusionPipeline.from_pretrained(
41
  "openart-custom/DynaVisionXL",
42
  torch_dtype=torch.float16,
43
  use_safetensors=True,
44
+ use_auth_token=HF_TOKEN # Inclui o token aqui
45
  )
46
  item_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(item_pipe.scheduler.config)
47
 
48
+ # Função de geração de imagem com alocação de GPU (através do decorador do Hugging Face Spaces)
49
+ @spaces.GPU(duration=60) # Aloca a GPU somente durante a execução desta função
50
  def generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
51
  if model_type == "character":
52
  pipe = character_pipe
53
  default_prompt = "1girl, souji okita, fate series, solo, upper body, bedroom, night, seducing, (sexy clothes)"
54
+ default_negative_prompt = ("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, "
55
+ "low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, "
56
+ "signature, extra digits, artistic error, username, scan, [abstract]")
57
  elif model_type == "item":
58
  pipe = item_pipe
59
  default_prompt = "great sword, runes on blade, acid on blade, weapon, (((item)))"
60
  default_negative_prompt = "1girl, girl, man, boy, 1man, men, girls"
61
  else:
62
+ return "Tipo inválido. Escolha entre 'character' ou 'item'."
63
 
64
+ # Se o usuário fornecer prompt, utiliza-o; caso contrário, usa o padrão
65
  final_prompt = prompt if prompt else default_prompt
66
  final_negative_prompt = negative_prompt if negative_prompt else default_negative_prompt
67
 
68
+ # Move o pipeline para a GPU
69
  pipe.to("cuda")
70
 
71
+ # Geração da imagem
72
+ result = pipe(
73
  prompt=final_prompt,
74
  negative_prompt=final_negative_prompt,
75
  width=int(width),
76
  height=int(height),
77
  guidance_scale=float(guidance_scale),
78
  num_inference_steps=int(num_inference_steps)
79
+ )
80
+ image = result.images[0]
81
 
82
+ # Salva a imagem em um arquivo temporário
83
  temp_file = "/tmp/generated_image.png"
84
  image.save(temp_file)
85
 
86
+ # Faz upload para o AWS S3
87
  file_name = datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
88
  try:
89
  s3_client.upload_file(temp_file, AWS_BUCKET_NAME, file_name)
90
  s3_url = f"https://{AWS_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{file_name}"
91
  return s3_url
92
  except NoCredentialsError:
93
+ return "Credenciais não disponíveis"
94
+
95
+ # Função que integra a geração via Gradio
96
+ def gradio_generate(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
97
+ return generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps)
98
+
99
+ # Definindo os componentes de entrada utilizando a API atual do Gradio
100
+ model_type_input = gr.Dropdown(choices=["character", "item"], value="character", label="Model Type")
101
+ prompt_input = gr.Textbox(lines=2, placeholder="Digite o prompt (deixe vazio para o padrão)", label="Prompt")
102
+ negative_prompt_input = gr.Textbox(lines=2, placeholder="Digite o negative prompt (deixe vazio para o padrão)", label="Negative Prompt")
103
+ width_input = gr.Number(value=512, label="Width")
104
+ height_input = gr.Number(value=512, label="Height")
105
+ guidance_scale_input = gr.Number(value=7.5, label="Guidance Scale")
106
+ num_inference_steps_input = gr.Number(value=50, label="Number of Inference Steps")
107
+
108
+ # Criação da interface Gradio
109
+ iface = gr.Interface(
110
+ fn=gradio_generate,
111
+ inputs=[
112
+ model_type_input,
113
+ prompt_input,
114
+ negative_prompt_input,
115
+ width_input,
116
+ height_input,
117
+ guidance_scale_input,
118
+ num_inference_steps_input,
119
+ ],
120
+ outputs="text",
121
+ title="Image Generation API",
122
+ description="Gere imagens usando modelos de difusão e faça upload para o AWS S3."
123
+ )
 
 
 
124
 
 
125
  if __name__ == "__main__":
126
+ iface.launch()
 
requirements.txt CHANGED
@@ -1,10 +1,6 @@
1
- fastapi
2
- uvicorn
3
- transformers
4
- spaces
5
  diffusers
6
  torch
7
  boto3
8
  python-dotenv
9
  Pillow
10
- accelerate
 
1
+ gradio
 
 
 
2
  diffusers
3
  torch
4
  boto3
5
  python-dotenv
6
  Pillow