Update app.py

app.py

@@ -1,118 +1,48 @@
-import gradio as gr
-from PIL import Image, ImageDraw, ImageFilter
 import torch
-from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline
-import torchvision.transforms as T
-from torchvision import models
-import numpy as np
+import gradio as gr
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from PIL import Image
 import cv2
+import numpy as np
 
-# Load the SDXL pipelines
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe_tshirt = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float32
-).to(device)
-
-pipe_design = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float32
-).to(device)
-
-# Generate T-shirt image
-def generate_cloth(color_prompt):
-    prompt = f"A plain {color_prompt} T-shirt hanging on a plain wall, realistic, high-quality photo."
-    image = pipe_tshirt(
-        prompt=prompt, strength=0.8, guidance_scale=7.5, num_inference_steps=30
-    ).images[0]
-    return image
-
-# Generate design image
-def generate_design(design_prompt):
-    prompt = f"A bold {design_prompt} design with vibrant colors, highly detailed."
-    image = pipe_design(
-        prompt=prompt, strength=0.8, guidance_scale=7.5, num_inference_steps=30
-    ).images[0]
-    return image.convert("RGBA")
-
-# Load pretrained DeepLabV3 model for T-shirt segmentation
-segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
-
-# Apply segmentation to extract T-shirt mask
-def get_tshirt_mask(image):
-    image = image.convert("RGB")  # Ensure 3 channels
-    preprocess = T.Compose([
-        T.Resize((520, 520)),  # Resize to avoid distortion
-        T.ToTensor(),
-        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    input_tensor = preprocess(image).unsqueeze(0)
-    with torch.no_grad():
-        output = segmentation_model(input_tensor)["out"][0]
-    
-    # Extract T-shirt mask (class 15 in COCO dataset)
-    mask = output.argmax(0).byte().cpu().numpy()
-    raw_mask = Image.fromarray((mask == 15).astype("uint8") * 255)  # Binary mask
-    processed_mask = post_process_mask(raw_mask)  # Apply post-processing
-    return processed_mask.resize(image.size)
-
-# Post-process mask to improve quality
-def post_process_mask(mask):
-    # Convert mask to NumPy array
-    mask_np = np.array(mask)
-    
-    # Morphological operations to refine mask
-    kernel = np.ones((5, 5), np.uint8)
-    mask_np = cv2.dilate(mask_np, kernel, iterations=2)  # Expand mask
-    mask_np = cv2.erode(mask_np, kernel, iterations=1)   # Remove noise
-    
-    # Convert back to PIL image and smooth
-    processed_mask = Image.fromarray(mask_np).filter(ImageFilter.GaussianBlur(3))
-    return processed_mask
-
-# Blend design seamlessly with T-shirt
-def blend_design_with_tshirt(cloth_image, design_image, mask):
-    # Resize the design to match T-shirt bounding box
-    mask_np = np.array(mask) / 255.0
-    cloth_np = np.array(cloth_image)
-    design_np = np.array(design_image.resize(cloth_image.size))
-
-    # Perform alpha blending with mask
-    blended = (design_np[..., :3] * mask_np[..., None] + cloth_np * (1 - mask_np[..., None])).astype(np.uint8)
-
-    return Image.fromarray(blended)
-
-# Main function to process T-shirt and design
-def design_tshirt(color_prompt, design_prompt):
-    try:
-        # Generate cloth and design images
-        cloth_image = generate_cloth(color_prompt)
-        design_image = generate_design(design_prompt)
-
-        # Get T-shirt mask
-        mask = get_tshirt_mask(cloth_image)
-
-        # Blend the design with the T-shirt
-        final_image = blend_design_with_tshirt(cloth_image, design_image, mask)
-        return final_image
-
-    except Exception as e:
-        return f"An error occurred: {str(e)}. Please try again with different inputs."
-
-# Gradio UI
-with gr.Blocks() as interface:
-    gr.Markdown("# **AI T-Shirt Designer**")
-    gr.Markdown("Generate custom T-shirts by specifying a color and adding a design that blends perfectly with the T-shirt.")
-    with gr.Row():
-        with gr.Column():
-            color_prompt = gr.Textbox(label="T-Shirt Color", placeholder="E.g., Red, Blue")
-            design_prompt = gr.Textbox(label="Design Details", placeholder="E.g., Abstract art, Nature patterns")
-            generate_button = gr.Button("Generate T-Shirt")
-        with gr.Column():
-            output_image = gr.Image(label="Final T-Shirt Design")
-
-    generate_button.click(
-        design_tshirt,
-        inputs=[color_prompt, design_prompt],
-        outputs=output_image,
-    )
-
-interface.launch(debug=True)
+# Load ControlNet pre-trained model for depth-based warping
+controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    torch_dtype=torch.float16
+).to("cuda")
+
+# Function to generate depth map of the cloth
+def generate_depth_map(image_path):
+    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+    depth_map = cv2.Laplacian(image, cv2.CV_64F)
+    depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+    return Image.fromarray(depth_map)
+
+# Function to blend design onto fabric using ControlNet
+def blend_design_on_cloth(fabric_image, design_image, prompt="T-shirt with embedded design"):
+    # Generate depth map for fabric
+    depth_map = generate_depth_map(fabric_image)
+
+    # Generate realistic blended output
+    result = pipe(prompt=prompt, image=fabric_image, control_image=depth_map, num_inference_steps=30).images[0]
+    return result
+
+# Gradio interface
+def process_image(fabric_image, design_image):
+    result = blend_design_on_cloth(fabric_image, design_image)
+    return result
+
+interface = gr.Interface(
+    fn=process_image,
+    inputs=[gr.Image(type="pil", label="Upload Cloth Image"),
+            gr.Image(type="pil", label="Upload Design/Text Image")],
+    outputs=gr.Image(type="pil", label="Blended Output"),
+    title="AI-Powered Text & Design Blending on Cloth",
+    description="Upload a T-shirt or fabric image and a design. AI will blend it naturally into the fabric!",
+)
+
+# Launch app for Hugging Face Spaces
+if __name__ == "__main__":
+    interface.launch()