main

app.py

@@ -0,0 +1,245 @@
+import os
+from io import BytesIO
+
+import gradio as gr
+import grpc
+from PIL import Image
+import pandas as pd
+
+import numpy as np
+import time
+
+from io import BytesIO
+from inference_pb2 import LoraRequest, LoraResponse
+from inference_pb2_grpc import LoraServiceStub
+import grpc
+
+
+
+PREFIX = "/home/jovyan/bobkov/lora_demo/DEMO/"
+
+info = {
+    'image': PREFIX + 'demo/{0}/{0}.jpg',
+    'weights_path': PREFIX + 'demo_results/flux-lora-{0}_aug-rank16',
+    'caption': PREFIX + 'demo/{0}_aug/data.csv',
+    'aug_path': PREFIX + 'demo/{0}_aug_filter/'
+}
+
+params = {
+    'cup': {'switch_t' : 7, 'aug_image' : None, 'checkpoint' : 1000},
+    'face_lifting': {'switch_t' : 7, 'aug_image' : None, 'checkpoint' : 1000},
+    'coffe_machine': {'switch_t' : 7, 'aug_image' : None, 'checkpoint' : 1000},
+    'kettle': {'switch_t' : 3, 'aug_image' : None, 'checkpoint' : 1000},
+    'body_lotion': {'switch_t' : 7, 'aug_image' : None, 'checkpoint' : 1000},
+    'toy': {'switch_t' : 3, 'aug_image' : None, 'checkpoint' : 1000},
+    'bag': {'switch_t' : 3, 'aug_image' : None, 'checkpoint' : 1000},
+    'armchair': {'switch_t' : 3, 'aug_image' : None, 'checkpoint' : 600},
+    'pendant': {'switch_t' : -1, 'aug_image' : None, 'checkpoint' : 1000},
+    'car': {'switch_t' : 7, 'aug_image' : 'car_aug_2.jpg', 'checkpoint' : 600},
+}
+
+table = pd.read_csv("/home/jovyan/bobkov/lora_demo/DEMO/demo/data.csv")
+CAPTIONS = {}
+for line in table.values:
+    CAPTIONS[line[0]] = line[1]
+    
+
+def bytes_to_image(image: bytes) -> Image.Image:
+    image = Image.open(BytesIO(image))
+    return image
+
+
+def generate_image(concept, prompt, progress=gr.Progress(track_tqdm=True)):
+    with grpc.insecure_channel(os.environ["SERVER"]) as channel:
+        stub = LoraServiceStub(channel)
+    
+        output = stub.generate(
+            LoraRequest(prompt=prompt, concept=concept, use_cache=False)
+        )
+        
+    return gr.update(value=bytes_to_image(output.res1)), gr.update(), gr.update(), gr.update()
+
+
+temaplte = """
+<div style="font-size: 18px;">
+    <b>Product description:</b> {}
+</div>
+"""
+
+
+def action1(): 
+    concept = "kettle" #############
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action2():
+    concept = "face_lifting" ###################
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action3():
+    concept = "pendant" #############
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action4():
+    concept = "car"
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action5():
+    concept = "body_lotion" #################
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action6():
+    concept = "toy" ############
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action7():
+    concept = "bag" ##############
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+def action8():
+    concept = "armchair" ############
+    img = Image.open(info["image"].format(concept))
+    description = temaplte.format(CAPTIONS[concept])
+    return gr.update(value=img, visible=True), gr.update(value=description, visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value=concept)
+
+
+css2 = """
+
+.my-custom-button {
+    width: 100px;          /* Button size */
+    height: 130px;
+    padding: 0;            /* Remove default padding */
+    margin: 0px;           /* Optional spacing between buttons */
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    background-color: transparent;
+    border: none;
+    overflow: hidden;      /* Ensures the image doesn't overflow */
+    --text-xl: 150px
+}
+
+.my-custom-button img {
+    max-width: 100%;
+    max-height: 100%;
+    object-fit: contain;   /* Ensure icon scales properly */
+}
+
+.input_image_container {
+    width: 350px !important;
+    height: 350px !important;
+    overflow: hidden;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    background-color: #f0f0f0;
+}
+.input_image_container img {
+    max-width: 100%;
+    max-height: 100%;
+    width: 350px;
+    height: 350px;
+    object-fit: contain;
+    display: block;
+    margin: 0 auto;
+}
+
+.prompt input {
+    font-size: 20px;
+}
+
+.prompt input::placeholder {
+    font-size: 20px;
+}
+
+.prompt label {
+    font-size: 20px !important;
+}
+"""
+
+def get_demo():
+    with gr.Blocks(css="""
+        .centered { 
+            display: flex; 
+            justify-content: center; 
+            align-items: center; 
+            height: 100%; 
+        }
+        .centered img {
+            margin: auto;
+            object-fit: contain;
+        }
+    """ + css2) as demo:
+        gr.Markdown("## Showcase Commercial Products with Stunning Natural Backgrounds")
+        with gr.Row():
+            with gr.Column(elem_classes=["centered"]):
+                with gr.Row():
+                    btn1 = gr.Button("", icon=info["image"].format("kettle"), elem_classes=["my-custom-button"])
+                    btn2 = gr.Button("", icon=info["image"].format("face_lifting"), elem_classes=["my-custom-button"])
+                    btn3 = gr.Button("", icon=info["image"].format("pendant"), elem_classes=["my-custom-button"])
+                    btn4 = gr.Button("", icon=info["image"].format("car"), elem_classes=["my-custom-button"])
+
+                with gr.Row():
+                    btn5 = gr.Button("", icon=info["image"].format("body_lotion"), elem_classes=["my-custom-button"])
+                    btn6 = gr.Button("", icon=info["image"].format("toy"), elem_classes=["my-custom-button"])
+                    btn7 = gr.Button("", icon=info["image"].format("bag"), elem_classes=["my-custom-button"])
+                    btn8 = gr.Button("", icon=info["image"].format("armchair"), elem_classes=["my-custom-button"])
+
+                prod_desc = gr.Markdown(value="""<div style="font-size: 20px;">Choose the product you want to showcase </div"> 🠕""", visible=True)
+
+                input_image = gr.Image(label="Chosen product", type="pil", height=300, width=300, visible=False, interactive=False, container=True, elem_classes=["input_image_container"])
+                descr = gr.Markdown(value=temaplte.format(""), visible=False)
+
+                concept = gr.Textbox("", visible=False)
+                prompt = gr.Textbox("", placeholder="is in the cozy kitchen", label="Describe the enviroment for your product", submit_btn=False, max_lines=1, visible=False, elem_classes=["prompt"])
+
+                btn_generate = gr.Button("Generate images", visible=False)
+
+
+                btn1.click(fn=action1, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn2.click(fn=action2, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn3.click(fn=action3, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn4.click(fn=action4, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn5.click(fn=action5, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn6.click(fn=action6, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn7.click(fn=action7, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+                btn8.click(fn=action8, outputs=[input_image, descr, prod_desc, prompt, btn_generate, concept])
+
+            with gr.Column():
+                with gr.Row():
+                    res1 = gr.Image(label="Result", visible=True)#, height=450, width=450)
+                    res2 = gr.Image(label="Result 2", visible=False)
+                with gr.Row():
+                    res3 = gr.Image(label="Result 2", visible=False, height=450, width=450)
+                    res4 = gr.Image(label="Result 4", visible=False) #<div style="text-align: center; font-size: 18px;">
+                gr.Markdown('''
+                    <div style="display: flex; align-items: center; gap: 10px; font-size: 20px; text-align: center; margin-left: 200px;">
+                        <div>Made by FusionBrainLab, AIRI</div><img src="https://static.tildacdn.com/tild3633-6662-4437-a333-646631346335/Airinet.png" style="width: 70px; height: auto;">
+                    </div>
+                    ''')
+
+        btn_generate.click(
+            fn=generate_image, 
+            inputs=[concept, prompt],
+            outputs=[res1, res2, res3, res4] # font-family: Arial, sans-serif;
+        )
+
+    return demo
+
+
+if __name__ == '__main__':
+    demo = get_demo()
+    demo.launch(server_name="0.0.0.0", server_port=7860)

inference_pb2.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: inference.proto
+# Protobuf Python Version: 6.31.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    0,
+    '',
+    'inference.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0finference.proto\x12\tinference\"A\n\x0bLoraRequest\x12\x0f\n\x07\x63oncept\x18\x01 \x01(\t\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x11\n\tuse_cache\x18\x03 \x01(\x08\"F\n\x0cLoraResponse\x12\x0c\n\x04res1\x18\x01 \x01(\x0c\x12\x0c\n\x04res2\x18\x02 \x01(\x0c\x12\x0c\n\x04res3\x18\x03 \x01(\x0c\x12\x0c\n\x04res4\x18\x04 \x01(\x0c\x32J\n\x0bLoraService\x12;\n\x08generate\x12\x16.inference.LoraRequest\x1a\x17.inference.LoraResponseb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'inference_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_LORAREQUEST']._serialized_start=30
+  _globals['_LORAREQUEST']._serialized_end=95
+  _globals['_LORARESPONSE']._serialized_start=97
+  _globals['_LORARESPONSE']._serialized_end=167
+  _globals['_LORASERVICE']._serialized_start=169
+  _globals['_LORASERVICE']._serialized_end=243
+# @@protoc_insertion_point(module_scope)

inference_pb2_grpc.py

@@ -0,0 +1,97 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import inference_pb2 as inference__pb2
+
+GRPC_GENERATED_VERSION = '1.73.1'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in inference_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class LoraServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.generate = channel.unary_unary(
+                '/inference.LoraService/generate',
+                request_serializer=inference__pb2.LoraRequest.SerializeToString,
+                response_deserializer=inference__pb2.LoraResponse.FromString,
+                _registered_method=True)
+
+
+class LoraServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def generate(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_LoraServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'generate': grpc.unary_unary_rpc_method_handler(
+                    servicer.generate,
+                    request_deserializer=inference__pb2.LoraRequest.FromString,
+                    response_serializer=inference__pb2.LoraResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'inference.LoraService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('inference.LoraService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class LoraService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def generate(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/inference.LoraService/generate',
+            inference__pb2.LoraRequest.SerializeToString,
+            inference__pb2.LoraResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)