app.py

import gradio as gr
import numpy as np
import random

import spaces #[uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline
import torch
import subprocess
from groq import Groq
import base64
import os

subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)

# Load FLUX image generator
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "black-forest-labs/FLUX.1-schnell"  # Replace to the model you would like to use
lora_path = "matteomarjanovic/flatsketcher"
weigths_file = "lora.safetensors"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)
pipe.load_lora_weights(lora_path, weight_name=weigths_file)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def encode_image(image_path):
  with open(image_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')


@spaces.GPU #[uncomment to use ZeroGPU]
def generate_description_fn(
    image,
    progress=gr.Progress(track_tqdm=True),
):
    base64_image = encode_image(image)

    client = Groq(
        api_key=os.environ.get("GROQ_API_KEY"),
    )

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": """
                            I want you to imagine how the black and white technical flat sketch of the garment you see in the picture would look like, and describe that flat sketch to me in rich details, both front and back descriptions are mandatory. The description of the flat sketch should be one paragraph long and you don't have to add any additional comment.

                            The style of the result should look somewhat like the following example:
                            The technical flat sketch of the dress would depict a midi-length, off-the-shoulder design with a smocked bodice and short puff sleeves that have elasticized cuffs. The elastic neckline sits straight across the chest and back, ensuring a secure fit. The bodice transitions into a flowy, tiered skirt with three evenly spaced gathered panels, creating soft volume. The back view mirrors the front, maintaining the smocked fit and tiered skirt without visible closures, suggesting a pullover style. Elasticized areas would be marked with textured lines, while the gathers and drape would be indicated through subtle curved strokes, ensuring clarity in construction details.
                            """
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        model="llama-3.2-11b-vision-preview",
    )

    prompt = chat_completion.choices[0].message.content + " In the style of FLTSKC"

    image = pipe(
        prompt=prompt,
        guidance_scale=0.,
        num_inference_steps=4,
        width=1420,
        height=1080,
        max_sequence_length=256,
        
    ).images[0]

    return prompt, image


css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

# generated_prompt = ""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# Draptic: from garment image to technical flat sketch")
    with gr.Row():
        with gr.Column(elem_id="col-input-image"):
            gr.Markdown(" ## Drop your image here")
            input_image = gr.Image(type="filepath")
        with gr.Column(elem_id="col-container"):
            generate_button = gr.Button("Generate flat sketch", scale=0, variant="primary")

            result = gr.Image(label="Result", show_label=False)

            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt = gr.Text(
                    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=1024,  # Replace with defaults that work for your model
                    )

                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,  # Replace with defaults that work for your model
                    )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=0.0,  # Replace with defaults that work for your model
                    )

                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=2,  # Replace with defaults that work for your model
                    )

            gr.Examples(examples=examples, inputs=[prompt])

        gr.on(
            triggers=[generate_button.click],
            fn=generate_description_fn,
            inputs=[
                input_image,
            ],
            outputs=[generated_prompt, result],
        )

if __name__ == "__main__":
    demo.launch()