app.py

from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline
import torch
import numpy as np
import roop.globals
from roop.core import (
    start,
    decode_execution_providers,
    suggest_max_memory,
    suggest_execution_threads,
)
from roop.processors.frame.core import get_frame_processors_modules
from roop.utilities import normalize_output_path
import os
from PIL import Image


base_model="stabilityai/stable-diffusion-xl-base-1.0"

# base = DiffusionPipeline.from_pretrained(base_model,
#         torch_dtype=torch.float16, use_safetensors=True
# ).to("cuda")
base = DiffusionPipeline.from_pretrained(base_model,
        torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
lora_model_path = "./pytorch_weights.safetensors"

base.load_lora_weights(lora_model_path)

def swap_face(source_file, target_file,doFaceEnhancer):

    source_path = "input.jpg"
    target_path = "target.jpg"

    # source_image = Image.fromarray(source_file)
    # source_image.save(source_path)
    # target_image = Image.fromarray(target_file)
    # target_image.save(target_path)
    
    # Open and save the source image
    source_image = Image.open(source_file)
    source_image.save(source_path)

    # Open and save the target image
    target_image = Image.open(target_file)
    target_image.save(target_path)

    print("source_path: ", source_path)
    print("target_path: ", target_path)

    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    output_path = "output.jpg"
    roop.globals.output_path = normalize_output_path(
        roop.globals.source_path, roop.globals.target_path, output_path
    )
    if doFaceEnhancer == True:
        roop.globals.frame_processors = ["face_swapper","face_enhancer"]
    else:
        roop.globals.frame_processors = ["face_swapper"]
    roop.globals.headless = True
    roop.globals.keep_fps = True
    roop.globals.keep_audio = True
    roop.globals.keep_frames = False
    roop.globals.many_faces = False
    roop.globals.video_encoder = "libx264"
    roop.globals.video_quality = 18
    roop.globals.max_memory = suggest_max_memory()
    roop.globals.execution_providers = decode_execution_providers(["cpu"])
    roop.globals.execution_threads = suggest_execution_threads()

    print(
        "start process",
        roop.globals.source_path,
        roop.globals.target_path,
        roop.globals.output_path,
    )

    for frame_processor in get_frame_processors_modules(
        roop.globals.frame_processors
    ):
        if not frame_processor.pre_check():
            return

    start()
    return output_path



def generate_image(prompt, height=1024, width=1024):
    
    image = base(
        prompt=prompt,
        negative_prompt="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime), text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
        num_inference_steps=40,
        guidance_scale=3,
        height=int(height),
        width=int(width),
        num_images_per_prompt=1,
    )
    

    for i in range(len(image.images)):
        # Save the generated images
        original_image_path = f"./output{i}.png"
        image.images[i].save(original_image_path)

        # The source image for face swapping
        source_image_path = "./roop/IMG_0991.png"

        # Run the face swap and optionally face enhancement
        swapped_image_path = swap_face(source_image_path, original_image_path, False)

        # Load the swapped image and add to the output list
        swapped_image = Image.open(swapped_image_path)

        return swapped_image

examples = [
    ["xkx man as superman. Uhd, 8k", 1024, 1024],
    ["photo of xkx man on a beach. Amazing scenery, waves, sun.  Uhd, 8k", 1024, 1536],
]
# Set up the Gradio interface
interface = gr.Interface(
    fn=generate_image,
inputs=[
        gr.Text(label="Prompt"),
        gr.Number(label="Height"),
        gr.Number(label="Width")
    ],
    outputs=gr.Image(label="Image"),
    examples=examples
            
)

interface.launch()