Update app.py

app.py

@@ -4,6 +4,8 @@ import random
 import spaces
 import torch
 from diffusers import DiffusionPipeline
+from PIL import Image
+import io
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -16,6 +18,41 @@ pipe = DiffusionPipeline.from_pretrained(
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+def create_tshirt_preview(design_image, tshirt_color="white"):
+    """
+    Overlay the design onto a t-shirt template
+    """
+    # Create a base t-shirt shape
+    tshirt_width = 800
+    tshirt_height = 1000
+    
+    # Create base t-shirt image
+    tshirt = Image.new('RGB', (tshirt_width, tshirt_height), tshirt_color)
+    
+    # Convert design to PIL Image if it's not already
+    if not isinstance(design_image, Image.Image):
+        design_image = Image.fromarray(design_image)
+    
+    # Resize design to fit nicely on shirt (30% of shirt width)
+    design_width = int(tshirt_width * 0.3)
+    design_height = int(design_width * design_image.size[1] / design_image.size[0])
+    design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
+    
+    # Calculate position to center design on shirt (top third of shirt)
+    x = (tshirt_width - design_width) // 2
+    y = int(tshirt_height * 0.25)  # Position in top third
+    
+    # If design has transparency (RGBA), create mask
+    if design_image.mode == 'RGBA':
+        mask = design_image.split()[3]
+    else:
+        mask = None
+    
+    # Paste design onto shirt
+    tshirt.paste(design_image, (x, y), mask)
+    
+    return tshirt
+
 def enhance_prompt_for_tshirt(prompt, style=None):
     """Add specific terms to ensure good t-shirt designs."""
     style_terms = {
@@ -27,9 +64,9 @@ def enhance_prompt_for_tshirt(prompt, style=None):
     }
     
     base_terms = [
-        "t-shirt design",
-        "centered composition",
-        "high quality",
+        "create a t-shirt design",
+        "with centered composition",
+        "4k high quality",
         "professional design",
         "clear background"
     ]
@@ -43,14 +80,17 @@ def enhance_prompt_for_tshirt(prompt, style=None):
     return enhanced_prompt
 
 @spaces.GPU()
-def infer(prompt, style=None, seed=42, randomize_seed=False, width=1024, height=1024, 
-          num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=False, 
+          width=1024, height=1024, num_inference_steps=4, 
+          progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
     enhanced_prompt = enhance_prompt_for_tshirt(prompt, style)
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
+    
+    # Generate the design
+    design_image = pipe(
         prompt=enhanced_prompt,
         width=width,
         height=height,
@@ -59,14 +99,24 @@ def infer(prompt, style=None, seed=42, randomize_seed=False, width=1024, height=
         guidance_scale=0.0
     ).images[0]
     
-    return image, seed
+    # Create t-shirt preview
+    tshirt_preview = create_tshirt_preview(design_image, tshirt_color)
+    
+    return design_image, tshirt_preview, seed
+
+# Available t-shirt colors
+TSHIRT_COLORS = {
+    "White": "#FFFFFF",
+    "Black": "#000000",
+    "Navy": "#000080",
+    "Gray": "#808080"
+}
 
 examples = [
-    ["Cool geometric mountain landscape", "minimal"],
-    ["Vintage motorcycle with flames", "vintage"],
-    ["Abstract watercolor butterfly", "artistic"],
-    ["Sacred geometry mandala", "geometric"],
-    ["Adventure Awaits typography", "typography"],
+    ["Cool geometric mountain landscape", "minimal", "White"],
+    ["Vintage motorcycle with flames", "vintage", "Black"],
+    ["Abstract watercolor butterfly in forest", "artistic", "White"],
+    ["Adventure Awaits typography", "typography", "Gray"]
 ]
 
 styles = [
@@ -105,13 +155,19 @@ css = """
     font-size: 1rem;
     transition: all 0.3s ease;
 }
+.results-row {
+    display: grid;
+    grid-template-columns: 1fr 1fr;
+    gap: 20px;
+    margin-top: 20px;
+}
 """
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(
             """
-            # 👕 T-Shirt Design Generator
+            # 👕 Deradh's T-Shirt Design Generator
             """,
             elem_classes=["main-title"]
         )
@@ -142,25 +198,30 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                     label="Style",
                     container=False
                 )
+            with gr.Column(scale=1):
+                tshirt_color = gr.Dropdown(
+                    choices=list(TSHIRT_COLORS.keys()),
+                    value="White",
+                    label="T-Shirt Color",
+                    container=False
+                )
             run_button = gr.Button(
                 "✨ Generate",
                 scale=0,
                 elem_classes=["generate-button"]
             )
         
-        with gr.Row():
-            with gr.Column():
-                result = gr.Image(
-                    label="Generated Design",
-                    show_label=True,
-                    elem_classes=["result-image"]
-                )
-            with gr.Column():
-                preview = gr.Image(
-                    label="T-Shirt Preview",
-                    show_label=True,
-                    elem_classes=["preview-image"]
-                )
+        with gr.Row(elem_classes=["results-row"]):
+            result = gr.Image(
+                label="Generated Design",
+                show_label=True,
+                elem_classes=["result-image"]
+            )
+            preview = gr.Image(
+                label="T-Shirt Preview",
+                show_label=True,
+                elem_classes=["preview-image"]
+            )
         
         with gr.Accordion("🔧 Advanced Settings", open=False):
             with gr.Group():
@@ -203,16 +264,16 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         gr.Examples(
             examples=examples,
             fn=infer,
-            inputs=[prompt, style],
-            outputs=[result, seed],
+            inputs=[prompt, style, tshirt_color],
+            outputs=[result, preview, seed],
             cache_examples=True
         )
         
         gr.on(
             triggers=[run_button.click, prompt.submit],
             fn=infer,
-            inputs=[prompt, style, seed, randomize_seed, width, height, num_inference_steps],
-            outputs=[result, seed]
+            inputs=[prompt, style, tshirt_color, seed, randomize_seed, width, height, num_inference_steps],
+            outputs=[result, preview, seed]
         )
 
 demo.launch()