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 from huggingface_hub import login 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, revision="fp16", torch_dtype=torch.float16, 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", revision="fp16", torch_dtype=torch.float16, 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)