helo2 / app.py
cngsm's picture
Update app.py
ef27658 verified
raw
history blame
1.25 kB
import os
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image
import gradio as gr
# Definir dispositivo de processamento
device = "cuda" if torch.cuda.is_available() else "cpu"
# Carregar o modelo Stable Diffusion Img2Img
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
torch_dtype=torch.float16
)
pipe = pipe.to(device)
# Função para gerar imagem
def generate_image(prompt, image):
# Processar a imagem com o modelo
result_image = pipe(
prompt=prompt,
image=image,
num_inference_steps=50,
strength=0.75,
guidance_scale=7.5
).images[0]
return result_image
# Configurar a interface Gradio
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt", placeholder="Digite o seu prompt aqui..."),
gr.Image(label="Imagem Inicial", type="pil") # Espera uma imagem no formato PIL
],
outputs=gr.Image(label="Imagem Gerada"),
title="Gerador de Imagens com Stable Diffusion Img2Img",
description="Insira um prompt e uma imagem inicial para gerar uma nova imagem."
)
# Executar a interface
if __name__ == "__main__":
iface.launch()