Update app.py

app.py

@@ -1,82 +1,97 @@
 import gradio as gr
-from PIL import Image
+import numpy as np
+import random
+from diffusers import DiffusionPipeline
 import torch
-from transformers import pipeline
-from torchvision import models, transforms
 
-# Load the models
-text_to_image_pipeline = pipeline("text-to-image-generation", model="CompVis/stable-diffusion-v1-4")
-segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True)
-segmentation_model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Define transformation for the segmentation model
-preprocess = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
+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)
 
-# Helper function to segment clothing area
-def segment_clothing(image):
-    input_tensor = preprocess(image)
-    input_batch = input_tensor.unsqueeze(0)
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
 
-    with torch.no_grad():
-        output = segmentation_model(input_batch)['out'][0]
-    output_predictions = output.argmax(0)
-    mask = output_predictions.byte().cpu().numpy()
-    
-    return mask
-
-# Function to generate base image
-def generate_base_image(base_prompt_part1, base_prompt_color, base_prompt_clothing):
-    # Combine the parts to create the full base prompt
-    base_prompt = f"{base_prompt_part1} {base_prompt_color} {base_prompt_clothing}"
+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}"
 
-    # Generate base clothing image
-    base_image = text_to_image_pipeline(base_prompt)[0]
-    base_image = Image.fromarray(base_image)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        
+    generator = torch.Generator().manual_seed(seed)
+    
+    try:
+        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]
+        print("Image generated successfully.")  # Debug: Confirm image generation
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return None
     
-    return base_image
+    return image
 
-# Define the function to generate design and paste it on the clothing
-def generate_and_paste_design(base_image, design_prompt):
-    # Generate design
-    generated_image = text_to_image_pipeline(design_prompt)[0]
-    generated_design = Image.fromarray(generated_image)
+examples = [
+    ["red", "t-shirt", "yellow stripes"],
+    ["blue", "hoodie", "minimalist"],
+    ["red", "sweatshirt", "geometric design"],
+]
 
-    # Segment the clothing area
-    clothing_mask = segment_clothing(base_image)
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
 
-    # Ensure the generated design fits within the clothing area
-    generated_design = generated_design.resize(base_image.size)
-    
-    # Paste the design onto the clothing area
-    clothing_area = Image.composite(generated_design, base_image, Image.fromarray(clothing_mask*255))
+power_device = "GPU" if torch.cuda.is_available() else "CPU"
 
-    return clothing_area
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""
+        # Text-to-Image Gradio Template
+        Currently running on {power_device}.
+        """)
+        
+        with gr.Row():
+            prompt_part1 = gr.Textbox(value="a single", label="Prompt Part 1", show_label=False, interactive=False, container=False, elem_id="prompt_part1", visible=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 grey 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)
 
-# Create the Gradio interface
-base_prompt_part1_input = gr.inputs.Textbox(lines=1, placeholder="Enter 'a single plain'")
-base_prompt_color_input = gr.inputs.Textbox(lines=1, placeholder="Enter color type")
-base_prompt_clothing_input = gr.inputs.Textbox(lines=1, placeholder="Enter clothing type")
-design_prompt_input = gr.inputs.Textbox(lines=1, placeholder="Enter design prompt")
-output_image = gr.outputs.Image(type="pil")
+        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=[prompt_part2, prompt_part3, prompt_part4])
 
-def full_process(base_prompt_part1, base_prompt_color, base_prompt_clothing, design_prompt):
-    # Generate the base image
-    base_image = generate_base_image(base_prompt_part1, base_prompt_color, base_prompt_clothing)
-    
-    # Generate and paste the design on the base image
-    final_image = generate_and_paste_design(base_image, design_prompt)
-    
-    return final_image
+    run_button.click(
+        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]
+    )
 
-gr.Interface(
-    fn=full_process,
-    inputs=[base_prompt_part1_input, base_prompt_color_input, base_prompt_clothing_input, design_prompt_input],
-    outputs=output_image,
-    title="Design and Paste on Clothing",
-    description="Generate a base clothing image from the given prompts and paste the generated design onto it."
-).launch()
+demo.queue().launch()