app.py

from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
from diffusers import StableDiffusionInpaintPipeline,StableDiffusionPipeline
from PIL import Image
import requests

import cv2
import torch
import matplotlib.pyplot as plt

import io
import requests

import os
import streamlit as st
from transformers import  AutoTokenizer, AutoModelForSeq2SeqLM, pipeline



processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
IPmodel_path = "runwayml/stable-diffusion-inpainting"

IPpipe = StableDiffusionInpaintPipeline.from_pretrained(
    IPmodel_path,
    use_auth_token= st.secrets["AUTH_TOKEN"]
).to(device)

trans_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
trans_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")


SDpipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token=st.secrets["AUTH_TOKEN"]).to(device)


def create_mask(image, prompt):
  inputs = processor(text=[prompt], images=[image], padding="max_length", return_tensors="pt")
  # predict
  with torch.no_grad():
    outputs = model(**inputs)

  preds = outputs.logits
  
  filename = f"mask.png"
  plt.imsave(filename,torch.sigmoid(preds))

  gray_image = cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)

  (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)

  # For debugging only:
  # cv2.imwrite(filename,bw_image)

  # fix color format
  cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)

  mask = cv2.bitwise_not(bw_image)
  cv2.imwrite(filename, mask)

  return Image.open('mask.png')
  
  


def generate_image(image, product_name, target_name):
  mask = create_mask(image, product_name)
  image = image.resize((512, 512))
  mask = mask.resize((512,512))
  guidance_scale=8
  #guidance_scale=16
  num_samples = 4

  prompt = target_name
  generator = torch.Generator(device=device).manual_seed(22) # change the seed to get different results

  im = IPpipe(
      prompt=prompt,
      image=image,
      mask_image=mask,
      guidance_scale=guidance_scale,
      generator=generator,
  ).images

  return im
  
  
  
def translate_sentence(article, source, target):
    if target == 'eng_Latn':
      return article
    translator = pipeline('translation', model=trans_model, tokenizer=trans_tokenizer, src_lang=source, tgt_lang=target)
    output = translator(article, max_length=400)
    output = output[0]['translation_text']
    return output


codes_as_string = '''Modern Standard Arabic	arb_Arab
Danish	dan_Latn
German	deu_Latn
Greek	ell_Grek
English	eng_Latn
Estonian	est_Latn
Finnish	fin_Latn
French	fra_Latn
Hebrew	heb_Hebr
Hindi	hin_Deva
Croatian	hrv_Latn
Hungarian	hun_Latn
Indonesian	ind_Latn
Icelandic	isl_Latn
Italian	ita_Latn
Japanese	jpn_Jpan
Korean	kor_Hang
Luxembourgish	ltz_Latn
Macedonian	mkd_Cyrl
Maltese	mlt_Latn
Dutch	nld_Latn
Norwegian Bokmål	nob_Latn
Polish	pol_Latn
Portuguese	por_Latn
Russian	rus_Cyrl
Slovak	slk_Latn
Slovenian	slv_Latn
Spanish	spa_Latn
Serbian	srp_Cyrl
Swedish	swe_Latn
Thai	tha_Thai
Turkish	tur_Latn
Ukrainian	ukr_Cyrl
Vietnamese	vie_Latn
Chinese (Simplified)	zho_Hans'''
    
codes_as_string = codes_as_string.split('\n')

flores_codes = {}
for code in codes_as_string:
    lang, lang_code = code.split('\t')
    flores_codes[lang] = lang_code
    
    
 
import gradio as gr
import gc 
gc.collect()

image_label = 'Please upload the image (optional)'
extract_label = 'Specify what need to be extracted from the above image'
prompt_label = 'Specify the description of image to be generated'
button_label = "Proceed"
output_label = "Generations"


shot_services = ['close-up', 'extreme-closeup', 'POV','medium', 'long']
shot_label = 'Choose the shot type'

style_services = ['polaroid', 'monochrome', 'long exposure','color splash', 'Tilt shift']
style_label = 'Choose the style type'

lighting_services = ['soft', 'ambivalent', 'ring','sun', 'cinematic']
lighting_label = 'Choose the lighting type'

context_services = ['indoor', 'outdoor', 'at night','in the park', 'in the beach','studio']
context_label = 'Choose the context'

lens_services = ['wide angle', 'telephoto', '24 mm','EF 70mm', 'Bokeh']
lens_label = 'Choose the lens type'

device_services = ['iphone', 'CCTV', 'Nikon ZFX','Canon', 'Gopro']
device_label = 'Choose the device type'


def change_lang(choice):
    global lang_choice 
    lang_choice = choice 
    new_image_label = translate_sentence(image_label, "english", choice)
    return [gr.update(visible=True, label=translate_sentence(image_label, flores_codes["English"],flores_codes[choice])), 
            gr.update(visible=True, label=translate_sentence(extract_label, flores_codes["English"],flores_codes[choice])), 
            gr.update(visible=True, label=translate_sentence(prompt_label, flores_codes["English"],flores_codes[choice])), 
            gr.update(visible=True, value=translate_sentence(button_label, flores_codes["English"],flores_codes[choice])),
            gr.update(visible=True, label=translate_sentence(button_label, flores_codes["English"],flores_codes[choice])), 
            ]

def add_to_prompt(prompt_text,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio ):
          if shot_radio != '':
            prompt_text += ","+shot_radio 
          if style_radio != '':
            prompt_text += ","+style_radio 
          if lighting_radio != '':
            prompt_text += ","+lighting_radio 
          if context_radio != '':
            prompt_text += ","+ context_radio 
          if lens_radio != '':
            prompt_text += ","+ lens_radio 
          if device_radio != '':
            prompt_text += ","+ device_radio 
          return prompt_text

def proceed_with_generation(input_file, extract_text, prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio):
    if extract_text == "" or input_file == "":
          translated_prompt = translate_sentence(prompt_text, flores_codes[lang_choice], flores_codes["English"])
          translated_prompt = add_to_prompt(translated_prompt,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
          print(translated_prompt)
          output = SDpipe(translated_prompt, height=512, width=512, num_images_per_prompt=4)
          return output.images
    elif extract_text != "" and input_file != "" and prompt_text !='':
          translated_prompt = translate_sentence(prompt_text, flores_codes[lang_choice], flores_codes["English"])
          translated_prompt = add_to_prompt(translated_prompt,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
          print(translated_prompt)
          translated_extract = translate_sentence(extract_text, flores_codes[lang_choice], flores_codes["English"])
          print(translated_extract)
          output = generate_image(Image.fromarray(input_file), translated_extract, translated_prompt)
          return output
    else:
          raise gr.Error("Please fill all details for guided image or atleast promt for free image rendition !")
          
    

with gr.Blocks() as demo:
                 
    lang_option =  gr.Dropdown(list(flores_codes.keys()), default='English', label='Please Select your Language')

    with gr.Row():
          input_file = gr.Image(interactive = True, label=image_label, visible=False, shape=(512,512))
          extract_text = gr.Textbox(label= extract_label, lines=1, interactive = True, visible = True)
          prompt_text = gr.Textbox(label= prompt_label, lines=1, interactive = True, visible = True)
    
    with gr.Accordion("Advanced Options", open=False):
          shot_radio = gr.Radio(shot_services  , label=shot_label, )
          style_radio = gr.Radio(style_services  , label=style_label)
          lighting_radio = gr.Radio(lighting_services  , label=lighting_label)
          context_radio = gr.Radio(context_services  , label=context_label)
          lens_radio = gr.Radio(lens_services  , label=lens_label)
          device_radio = gr.Radio(device_services  , label=device_label)

    button = gr.Button(value = button_label , visible = False)
   
    with gr.Row():
         output_gallery = gr.Gallery(label = output_label, visible= False)

    
    

    lang_option.change(fn=change_lang, inputs=lang_option, outputs=[input_file, extract_text, prompt_text, button, output_gallery])
    button.click( proceed_with_generation,  [input_file, extract_text, prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio], [output_gallery])
    

    demo.launch(debug=True)