Update app.py

app.py

@@ -14,6 +14,9 @@ import os
 import streamlit as st
 from transformers import  AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
+import gradio as gr
+import gc 
+
 
 
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
@@ -83,74 +86,17 @@ def generate_image(image, product_name, target_name):
       generator=generator,
   ).images
 
+  im.enable_attention_slicing()
   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()
+torch.cuda.empty_cache()
 
 image_label = 'Please upload the image (optional)'
-extract_label = 'Specify what need to be extracted from the above image'
+extract_label = 'Specify what needs to be extracted from the above image'
 prompt_label = 'Specify the description of image to be generated'
 button_label = "Proceed"
-output_label = "Generations"
-
+output_label = "Results"
 
 shot_services = ['close-up', 'extreme-closeup', 'POV','medium', 'long']
 shot_label = 'Choose the shot type'
@@ -171,56 +117,40 @@ 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 
+            prompt_text += ", "+shot_radio 
           if style_radio != '':
-            prompt_text += ","+style_radio 
+            prompt_text += ", "+style_radio 
           if lighting_radio != '':
-            prompt_text += ","+lighting_radio 
+            prompt_text += ", "+lighting_radio 
           if context_radio != '':
-            prompt_text += ","+ context_radio 
+            prompt_text += ", "+ context_radio 
           if lens_radio != '':
-            prompt_text += ","+ lens_radio 
+            prompt_text += ", "+ lens_radio 
           if device_radio != '':
-            prompt_text += ","+ 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)
+          prompt = add_to_prompt(prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
+          print(prompt)
+          output = SDpipe(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)
+          prompt = add_to_prompt(prompt_text,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
+          print(prompt)
+          print(extract_text)
+          output = generate_image(Image.fromarray(input_file), extract_text, prompt)
           return output
     else:
-          raise gr.Error("Please fill all details for guided image or atleast promt for free image rendition !")
+          raise gr.Error("Please fill all details for guided image or atleast prompt 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)
@@ -234,15 +164,12 @@ with gr.Blocks() as demo:
           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)
+    button = gr.Button(value = button_label , visible = True)
    
     with gr.Row():
-         output_gallery = gr.Gallery(label = output_label, visible= False)
+         output_gallery = gr.Gallery(label = output_label, visible= True)
 
-    
-    
 
-    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])