Spaces:
Running
Running
File size: 3,399 Bytes
21e64dd bac49e6 c9beab1 21e64dd c9beab1 21e64dd c9beab1 21e64dd c9beab1 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import gradio as gr
import torch
from PIL import Image
import numpy as np
import cv2
from diffusers import StableDiffusionPipeline
# Setup the model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "s3nh/artwork-arcane-stable-diffusion"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe = pipe.to(device)
# Generate T-shirt design function
def generate_tshirt_design(text):
prompt = f"{text}"
image = pipe(prompt).images[0]
return image
# Remove background from the generated design
def remove_background(design_image):
design_np = np.array(design_image)
gray = cv2.cvtColor(design_np, cv2.COLOR_BGR2GRAY)
_, alpha = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
b, g, r = cv2.split(design_np)
rgba = [b, g, r, alpha]
design_np = cv2.merge(rgba, 4)
return design_np
# Blend design with T-shirt mockup
def blend_design_with_mockup(mockup, design):
mockup_np = np.array(mockup)
design_np = np.array(design)
# Ensure the design is RGBA
if design_np.shape[2] == 3:
alpha = np.ones((design_np.shape[0], design_np.shape[1], 1), dtype=design_np.dtype) * 255
design_np = np.concatenate((design_np, alpha), axis=2)
design_resized = cv2.resize(design_np, (mockup_np.shape[1] // 4, mockup_np.shape[0] // 4)) # Adjust size as needed
y_offset = (mockup_np.shape[0] - design_resized.shape[0]) // 2
x_offset = (mockup_np.shape[1] - design_resized.shape[1]) // 2
y1, y2 = y_offset, y_offset + design_resized.shape[0]
x1, x2 = x_offset, x_offset + design_resized.shape[1]
alpha_s = design_resized[:, :, 3] / 255.0
alpha_l = 1.0 - alpha_s
for c in range(0, 3):
mockup_np[y1:y2, x1:x2, c] = (alpha_s * design_resized[:, :, c] +
alpha_l * mockup_np[y1:y2, x1:x2, c])
result_image = Image.fromarray(mockup_np)
return result_image
# T-shirt mockup generator with Gradio interface
examples = [
["MyBrand"],
["Hello World"],
["Team logo"],
]
css = """
#col-container {
margin: 0 auto;
max-width: 520px;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown("""
# T-shirt Design Generator with Stable Diffusion
""")
with gr.Row():
text = gr.Textbox(
label="Text",
placeholder="Enter text for the T-shirt design",
visible=True,
)
run_button = gr.Button("Generate Design", scale=0)
result = gr.Image(label="Design", show_label=False)
gr.Examples(
examples=examples,
inputs=[text]
)
def generate_tshirt_mockup(text):
# Generate T-shirt design
design_image = generate_tshirt_design(text)
# Remove background from design image
design_np = remove_background(design_image)
# Load blank T-shirt mockup template image
mockup_template = Image.open("/content/drive/MyDrive/unnamed.jpg")
# Blend design with mockup
result_image = blend_design_with_mockup(mockup_template, Image.fromarray(design_np))
return result_image
run_button.click(
fn=generate_tshirt_mockup,
inputs=[text],
outputs=[result]
)
demo.queue().launch()
|