Update app.py

app.py

@@ -1,65 +1,47 @@
 import gradio as gr
-import torch
-from PIL import Image
 import numpy as np
-import cv2
-from diffusers import StableDiffusionPipeline
+import random
+from diffusers import DiffusionPipeline
+import torch
 
-# 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]
+if torch.cuda.is_available():
+    torch.cuda.max_memory_allocated(device=device)
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to(device)
+else: 
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
+    pipe = pipe.to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def infer(prompt_part1, color, dress_type, design, prompt_part5, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    prompt = f"{prompt_part1} {color} colored plain {dress_type} with {design} design, {prompt_part5}"
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        
+    generator = torch.Generator().manual_seed(seed)
+    
+    image = pipe(
+        prompt=prompt, 
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale, 
+        num_inference_steps=num_inference_steps, 
+        width=width, 
+        height=height,
+        generator=generator
+    ).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"],
+    "red, t-shirt, yellow stripes",
+    "blue, hoodie, minimalist",
+    "red, sweat shirt, geometric design",
 ]
 
 css = """
@@ -69,47 +51,134 @@ css = """
 }
 """
 
+if torch.cuda.is_available():
+    power_device = "GPU"
+else:
+    power_device = "CPU"
+
 with gr.Blocks(css=css) as demo:
+    
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("""
-        # T-shirt Design Generator with Stable Diffusion
+        gr.Markdown(f"""
+        # Text-to-Image Gradio Template
+        Currently running on {power_device}.
         """)
-
+        
         with gr.Row():
-            text = gr.Textbox(
-                label="Text",
-                placeholder="Enter text for the T-shirt design",
-                visible=True,
+            
+            prompt_part1 = gr.Textbox(
+                value="a single", 
+                label="Prompt Part 1",
+                show_label=False,
+                interactive=False,
+                container=False,
+                elem_id="prompt_part1",
+                visible=False,
             )
-
-            run_button = gr.Button("Generate Design", scale=0)
-
-        result = gr.Image(label="Design", show_label=False)
-
+            
+            prompt_part2 = gr.Textbox(
+                label="color",
+                show_label=False,
+                max_lines=1,
+                placeholder="color (e.g., color category)",
+                container=False,
+            )
+            
+            prompt_part3 = gr.Textbox(
+                label="dress_type",
+                show_label=False,
+                max_lines=1,
+                placeholder="dress_type (e.g., t-shirt, sweatshirt, shirt, hoodie)",
+                container=False,
+            )
+            
+            prompt_part4 = gr.Textbox(
+                label="design",
+                show_label=False,
+                max_lines=1,
+                placeholder="design",
+                container=False,
+            )
+            
+            prompt_part5 = gr.Textbox(
+                value="hanging on the plain wall", 
+                label="Prompt Part 5",
+                show_label=False,
+                interactive=False,
+                container=False,
+                elem_id="prompt_part5",
+                visible=False,
+            )
+            
+            
+            run_button = gr.Button("Run", scale=0)
+        
+        result = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+            
+            negative_prompt = gr.Textbox(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
+            )
+            
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            
+            with gr.Row():
+                
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
+                )
+                
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
+                )
+            
+            with gr.Row():
+                
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                )
+                
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=12,
+                    step=1,
+                    value=2,
+                )
+        
         gr.Examples(
             examples=examples,
-            inputs=[text]
+            inputs=[prompt_part2]
         )
 
-    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("OIP (3).jpeg")
-
-        # 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],
+        fn=infer,
+        inputs=[prompt_part1, prompt_part2, prompt_part3, prompt_part4, prompt_part5, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs=[result]
     )
 
-demo.queue().launch()
+demo.queue().launch()