app.py

import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch
from PIL import Image, ImageOps, ImageEnhance

device = "cuda" if torch.cuda.is_available() else "cpu"

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, front_design, back_design, prompt_part5, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    front_prompt = f"front view of {prompt_part1} {color} colored plain {dress_type} with {front_design} design, {prompt_part5}"
    back_prompt = f"back view of {prompt_part1} {color} colored plain {dress_type} with {back_design} design, {prompt_part5}"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    generator = torch.Generator().manual_seed(seed)
    
    front_image = pipe(
        prompt=front_prompt, 
        guidance_scale=guidance_scale, 
        num_inference_steps=num_inference_steps, 
        width=width, 
        height=height,
        generator=generator
    ).images[0]
    
    back_image = pipe(
        prompt=back_prompt, 
        guidance_scale=guidance_scale, 
        num_inference_steps=num_inference_steps, 
        width=width, 
        height=height,
        generator=generator
    ).images[0]
    
    return front_image, back_image

def edit_image(img_data, operation, *args):
    image = Image.open(img_data)
    
    if operation == "rotate":
        angle = int(args[0])
        image = image.rotate(angle, expand=True)
    elif operation == "crop":
        left, top, right, bottom = map(int, args)
        image = image.crop((left, top, right, bottom))
    elif operation == "resize":
        width, height = map(int, args)
        image = image.resize((width, height))
    elif operation == "flip":
        if args[0] == "horizontal":
            image = ImageOps.mirror(image)
        else:
            image = ImageOps.flip(image)
    elif operation == "color":
        factor = float(args[0])
        image = ImageEnhance.Color(image).enhance(factor)
    
    return image

examples = [
    ["red", "t-shirt", "yellow stripes", "polka dots"],
    ["blue", "hoodie", "minimalist", "abstract art"],
    ["red", "sweat shirt", "geometric design", "plain"],
]

if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown(f"""
        # GenZ Couture
        Currently running on {power_device}.
        """)

    with gr.Row():
        with gr.Column():
            prompt_part1 = gr.Textbox(value="a single", label="Prompt Part 1")
            prompt_part2 = gr.Textbox(label="Color", placeholder="Color (e.g., red, blue)")
            prompt_part3 = gr.Textbox(label="Dress Type", placeholder="Dress Type (e.g., t-shirt, hoodie)")
            prompt_part4_front = gr.Textbox(label="Front Design", placeholder="Front Design")
            prompt_part4_back = gr.Textbox(label="Back Design", placeholder="Back Design")
            prompt_part5 = gr.Textbox(value="hanging on the plain wall", label="Prompt Part 5")
            seed = gr.Slider(0, MAX_SEED, step=1, label="Seed", value=42)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            width = gr.Slider(256, MAX_IMAGE_SIZE, step=32, label="Width", value=512)
            height = gr.Slider(256, MAX_IMAGE_SIZE, step=32, label="Height", value=512)
            guidance_scale = gr.Slider(1, 20, step=0.5, label="Guidance Scale", value=7.5)
            num_inference_steps = gr.Slider(10, 100, step=1, label="Number of Inference Steps", value=50)
            
            run_button = gr.Button("Generate Designs")
        
        with gr.Column():
            front_result = gr.Image(label="Front View Result", type="pil", interactive=True)
            back_result = gr.Image(label="Back View Result", type="pil", interactive=True)

    run_button.click(
        fn=infer,
        inputs=[prompt_part1, prompt_part2, prompt_part3, prompt_part4_front, prompt_part4_back, prompt_part5, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[front_result, back_result]
    )

    gr.Markdown("## Creative Touch")

    with gr.Row():
        edit_operation = gr.Dropdown(choices=["rotate", "crop", "resize", "flip", "color"], label="Edit Operation")
        edit_args = gr.Textbox(label="Edit Arguments (comma-separated)", placeholder="For rotate: angle, For crop: left,top,right,bottom, For resize: width,height, For flip: horizontal/vertical, For color: factor")

        edit_button = gr.Button("Edit Front Design")
        edited_image = gr.Image(label="Edited Front Design", type="pil", interactive=True)

        edit_button.click(
            fn=lambda img_data, operation, args: edit_image(img_data, operation, *args.split(',')),
            inputs=[front_result, edit_operation, edit_args],
            outputs=[edited_image]
        )

demo.queue().launch()