Update app.py

app.py

@@ -4,14 +4,11 @@ from PIL import Image
 import numpy as np
 import cv2
 from diffusers import StableDiffusionPipeline
-from huggingface_hub import login
-
-
 
 # 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, use_auth_token=True)
+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
@@ -30,6 +27,34 @@ def remove_background(design_image):
     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"],
@@ -76,28 +101,8 @@ with gr.Blocks(css=css) as demo:
         # Load blank T-shirt mockup template image
         mockup_template = Image.open("/content/drive/MyDrive/unnamed.jpg")
 
-        # Convert mockup template to numpy array
-        mockup_np = np.array(mockup_template)
-
-        # Resize design image to fit mockup
-        design_resized = cv2.resize(design_np, (mockup_np.shape[1] // 4, mockup_np.shape[0] // 4))  # Adjust size as needed
-
-        # Center the design on the mockup
-        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]
-
-        # Blend design with mockup using alpha channel
-        alpha_s = design_resized[:, :, 3] / 255.0 if design_resized.shape[2] == 4 else np.ones(design_resized.shape[:2])
-        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])
-
-        # Convert back to PIL image for Gradio output
-        result_image = Image.fromarray(mockup_np)
+        # Blend design with mockup
+        result_image = blend_design_with_mockup(mockup_template, Image.fromarray(design_np))
 
         return result_image
 
@@ -107,4 +112,4 @@ with gr.Blocks(css=css) as demo:
         outputs=[result]
     )
 
-demo.queue().launch()
+demo.queue().launch()