|
import gradio as gr |
|
from PIL import Image |
|
import torch |
|
from diffusers import StableDiffusionImg2ImgPipeline |
|
import spaces |
|
|
|
|
|
@spaces.GPU(duration=120) |
|
def initialize_pipeline(): |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model_id = "CompVis/stable-diffusion-v1-4" |
|
try: |
|
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id).to(device) |
|
except ValueError as e: |
|
return f"Error cargando el pipeline: {e}" |
|
return pipe |
|
|
|
|
|
pipe = None |
|
|
|
def setup_pipeline(): |
|
global pipe |
|
if pipe is None: |
|
pipe = initialize_pipeline() |
|
|
|
def cartoonize_image(image, prompt): |
|
setup_pipeline() |
|
|
|
try: |
|
|
|
image = image.convert("RGB") |
|
image = image.resize((512, 512)) |
|
|
|
|
|
with torch.no_grad(): |
|
result = pipe(prompt=prompt, init_image=image, strength=0.75).images[0] |
|
|
|
return result |
|
except Exception as e: |
|
return f"Error procesando la imagen: {e}" |
|
|
|
def main(): |
|
with gr.Blocks() as demo: |
|
gr.Markdown("## Conversi贸n de Imagen a Caricatura Estilo Pixar") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image(type="pil", label="Sube tu imagen") |
|
prompt = gr.Textbox(label="Prompt", value="A Pixar-style cartoon character") |
|
|
|
with gr.Row(): |
|
submit_button = gr.Button("Submit") |
|
stop_button = gr.Button("Stop") |
|
|
|
output_image = gr.Image(type="pil", label="Imagen Caricaturizada") |
|
|
|
def process_image(image, prompt): |
|
return cartoonize_image(image, prompt) |
|
|
|
submit_button.click( |
|
fn=process_image, |
|
inputs=[input_image, prompt], |
|
outputs=output_image |
|
) |
|
|
|
|
|
|
|
stop_button.click( |
|
fn=lambda: "El proceso ha sido detenido", |
|
outputs=output_image |
|
) |
|
|
|
demo.launch() |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|