Spaces:
Running
Running
File size: 2,657 Bytes
3bbb2bc 5a1023e 0fa0e0e 53f544e 0fa0e0e 5a1023e 53f544e 0fa0e0e 5a1023e 0fa0e0e 53f544e 5a1023e 0fa0e0e 5a1023e 3bbb2bc 0fa0e0e 5a1023e 0fa0e0e 5a1023e 53f544e 3bbb2bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import numpy as np
import cv2
# Device setup
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load Stable Diffusion model
model_name = "CompVis/stable-diffusion-v1-4"
def load_image_model():
return StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
image_model = load_image_model()
def generate_image(prompt):
with torch.no_grad():
image = image_model(prompt).images[0]
return image
# Face swap function placeholder (replace with actual implementation)
def swap_faces(image1, image2):
# Placeholder implementation, you should replace this with your actual face swapping code
return image1
def face_swap(image1, image2):
image1 = cv2.cvtColor(np.array(image1), cv2.COLOR_RGB2BGR)
image2 = cv2.cvtColor(np.array(image2), cv2.COLOR_RGB2BGR)
swapped_image = swap_faces(image1, image2)
return cv2.cvtColor(swapped_image, cv2.COLOR_BGR2RGB)
# Upscaling function
def upscale_image(image, scale_factor=2):
image = Image.fromarray(np.array(image))
width, height = image.size
new_size = (int(width * scale_factor), int(height * scale_factor))
upscaled_image = image.resize(new_size, Image.LANCZOS)
return upscaled_image
# Gradio interface function
def process_image(prompt, image1=None, image2=None, scale_factor=2):
# Generate image from prompt
generated_image = generate_image(prompt)
# Perform faceswap if two images are provided
swapped_image = None
if image1 and image2:
swapped_image = face_swap(image1, image2)
# Perform upscaling if scale factor is specified
upscaled_image = upscale_image(generated_image, scale_factor)
return {
"Generated Image": generated_image,
"Swapped Image": swapped_image,
"Upscaled Image": upscaled_image
}
# Gradio interface setup
iface = gr.Interface(
fn=process_image,
inputs=[
gr.Textbox(label="Enter your prompt"),
gr.Image(label="Image 1 (for faceswap)", type="pil", optional=True),
gr.Image(label="Image 2 (for faceswap)", type="pil", optional=True),
gr.Slider(label="Upscale Factor", minimum=1, maximum=4, step=1, default=2)
],
outputs=[
gr.Image(label="Generated Image"),
gr.Image(label="Swapped Image"),
gr.Image(label="Upscaled Image")
],
title="Fooocus Image Processing",
description="Generate images from prompts, swap faces, and upscale images."
)
if __name__ == "__main__":
iface.launch(server_name="0.0.0.0", server_port=7860)
|