# -*- coding: utf-8 -*-
"""Untitled11.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/17mNpy3A7oaJ-q9Txp8f4eBvRfO_k7NCb
"""

!pip install diffusers
!pip install transformers

import torch
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image
import gradio as gr

# Dispositivo de processamento
device = "cuda"

# Carregar o modelo CompVis/stable-diffusion-v1-4 (compatível com Img2ImgPipeline)
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",  # Modelo compatível
    torch_dtype=torch.float16
)
pipe = pipe.to(device)

# Carregar a imagem inicial
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
init_image = load_image(url).resize((1024, 1024))

# Prompt a ser usado para geração
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"

# Função de geração de imagem a partir de um prompt e uma imagem inicial
def generate_image(prompt, image):
    image = pipe(
        prompt=prompt,
        image=image,
        num_inference_steps=50,  # Aumentar o número de passos pode melhorar a qualidade
        strength=0.75,  # Reduzir a força ajuda a manter mais da imagem original (rosto consistente)
        guidance_scale=7.5  # Controle de qualidade para manter o rosto mais detalhado
    ).images[0]
    return image

# Interface com Gradio para facilitar a geração de imagens
gr.Interface(
    fn=generate_image,
    inputs=["text", "image"],
    outputs="image"
).launch()