Added Style embeddings

app.py

@@ -1,34 +1,48 @@
 import gradio as gr
 from torchvision import transforms
 import torch
-from main_inference import generate_mixed_image, generate_image, progress_video
+from main_inference import generate_mixed_image, generate_image
+from style_guidence import generate_with_prompt_style
 import matplotlib.colors as mcolors
 
-def run_generate_mixed_image(prompt1,prompt2):
-    image = generate_mixed_image(prompt1,prompt2)
+
+style_file_maps = {
+                   '3D Female Cyborgs':"style_embeddings/3d_female_cyborgs.bin", 
+                   '80s Anime':"style_embeddings/80s_anime.bin",
+                   'Anders Zorn':"style_embeddings/anders_zorn.bin",
+                   "Angus Mcbride":"style_embeddings/angus_mcbride.bin",
+                   "Breack Core":"style_embeddings/breakcore.bin", 
+                   "Brittney Williams":"style_embeddings/brittney_williams.bin",
+                   "Bull vs Bear":"style_embeddings/bull_vs_bear.bin",
+                   "Caitlin FairChild":"style_embeddings/caitlin_fairchild.bin",
+                   "Exodus Styling":"style_embeddings/exodus_styling.bin",
+                   "FoorByv2":"style_embeddings/foorbyv2.bin"
+                   }
+
+def run_generate_mixed_image(prompt1,prompt2,num_of_inf_steps,noised_image):
+    image = generate_mixed_image(prompt1,prompt2,num_of_inf_steps,noised_image)
     return image
 
-def run_generate_image(prompt1,noise_checkbox):
-    image = generate_image(prompt1,noised_image=noise_checkbox)
+def run_generate_image(prompt1,num_of_inf_steps,noise_checkbox):
+    image = generate_image(prompt1,num_inference_steps=num_of_inf_steps,noised_image=noise_checkbox)
     return image
 
-def run_generate_image_with_color_doninance(prompt1,color,color_loss_scale,noised_image_checkbox_1):
+def run_generate_image_with_color_doninance(prompt1,color,color_loss_scale,num_of_inf_steps,noised_image_checkbox_1):
     # Convert the hexadecimal color code to RGB values
     rgba_color = mcolors.hex2color(color)
     # Multiply the RGB values by 255 to get them in the [0, 255] range
     rgb_values = [int(val * 255) for val in rgba_color]
-    image = generate_image(prompt1,True,color,color_loss_scale,noised_image_checkbox_1)
+    image = generate_image(prompt1,num_of_inf_steps,True,rgb_values,color_loss_scale,noised_image_checkbox_1)
     return image
-    
-def run_process_video(prompt):
-    # Ask for text input
-    video = progress_video(prompt)
-    return video
+
+def run_generate_image_with_style(prompt,style,num_of_inf_steps):
+    output = generate_with_prompt_style(prompt, style_file_maps[style],num_of_inf_steps)
+    return output 
 
 
 description_text_to_image = """ ### Text to Image Generation
                 
-                1. Write a Text Prompt.
+                1. Write a Text Prompt and number of inference steps, the more the better results but execution time will be high.
                 
                 2. Output will be an image based on the text prompt provided.
                 
@@ -38,7 +52,7 @@ description_text_to_image = """ ### Text to Image Generation
 
 description_generate_mixed_image = """ ### Mix Image Generation
                 
-                1. Write Two Text prompts.
+                1. Write Two Text prompts and number of inference steps, the more the better results but execution time will be high.
                 
                 2. Output will a image which is mix of both of the text provided.
                 
@@ -48,7 +62,7 @@ description_generate_mixed_image = """ ### Mix Image Generation
 
 description_generate_image_with_color_dominance = """ ### Generate Images with color dominance
                 
-                1. Write a Text Prompt.
+                1. Write a Text Prompt and number of inference steps, the more the better results but execution time will be high.
                 
                 2. Select a color 
                 
@@ -60,11 +74,13 @@ description_generate_image_with_color_dominance = """ ### Generate Images with c
                                      
               """
               
-description_progress_video = """ ### Get the full generation process video
+description_generate_prompt_with_style = """ ### Get a generated image in the selection of your style
                 
-                1. Write a Text Prompt.
+                1. Write a Text Prompt and number of inference steps, the more the better results but execution time will be high.
                 
-                2. Output will be the video which contains frames of generated image, during various inference steps
+                2. Select a style to dominate the photo
+                
+                3. Get the Output
                                      
               """
 
@@ -75,18 +91,26 @@ image_input1 = gr.Image(type='pil')
 image_input2 = gr.Image(type='filepath') 
 image_input3 = gr.Image(type='pil')
 image_input4 = gr.Image(type='pil')
+image_input5 = gr.Image(type='pil')
 text_input = gr.Text(label="Enter Text Prompt")
 text_input2 = gr.Text(label="Enter Text Prompt")
 text_input3 = gr.Text(label="Enter Text Prompt")
 text_input4 = gr.Text(label = "Enter Text Prompt")
 text_input5 = gr.Text(label = "Enter Text Prompt")
-video_output = gr.Video()
+
+num_of_inf_steps_slider1 = gr.inputs.Slider(minimum=0, maximum=50, default=30, step=1,label="Num of Inference Steps")
+num_of_inf_steps_slider2 = gr.inputs.Slider(minimum=0, maximum=50, default=30, step=1,label="Num of Inference Steps")
+num_of_inf_steps_slider3 = gr.inputs.Slider(minimum=0, maximum=50, default=30, step=1,label="Num of Inference Steps")
+num_of_inf_steps_slider4 = gr.inputs.Slider(minimum=0, maximum=50, default=30, step=1,label="Num of Inference Steps")
 
 color = gr.ColorPicker(label="Select a Color",description="Choose a color from the color picker:")
 noised_image_checkbox  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
 noised_image_checkbox_1  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
 noised_image_checkbox_2  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
+noised_image_checkbox_3  = gr.inputs.Checkbox(default=False, label="Show Noised Image")
 color_loss_scale = gr.inputs.Slider(minimum=0, maximum=255, default=40, step=1,label="Color Loss")
+style_options = ['3D Female Cyborgs', '80s Anime','Anders Zorn',"Angus Mcbride","Breack Core", "Brittney Williams","Bull vs Bear","Caitlin FairChild","Exodus Styling","FoorByv2"]
+selected_style = gr.Dropdown(style_options,label="Select a Style to Follow",default="Anders Zorn")
 css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
 
 with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
@@ -100,6 +124,7 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
                 text_input.render()
+                num_of_inf_steps_slider1.render()
                 noised_image_checkbox.render()
             with gr.Column(scale=1):
                 image_input1.render()
@@ -118,7 +143,7 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
                 #             examples_per_page=4)
 
     run_generate_image_button.click(run_generate_image,
-                        inputs=[text_input,noised_image_checkbox],
+                        inputs=[text_input,num_of_inf_steps_slider1,noised_image_checkbox],
                         outputs=image_input1)
 
     with gr.Tab("Generate Image with Color Dominance"):
@@ -127,6 +152,7 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
             with gr.Column(scale=1):
                 text_input4.render()
                 color_loss_scale.render()
+                num_of_inf_steps_slider2.render()
                 noised_image_checkbox_1.render()
                 color.render()
             with gr.Column(scale=1):
@@ -146,7 +172,7 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
                 #             examples_per_page=4)
 
     run_generate_image_with_color_doninance_button.click(run_generate_image_with_color_doninance,
-                        inputs=[text_input4,color,color_loss_scale,noised_image_checkbox_1],
+                        inputs=[text_input4,color,color_loss_scale,num_of_inf_steps_slider2,noised_image_checkbox_1],
                         outputs=image_input3)
 
     ####################################################################################################################
@@ -156,6 +182,7 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
             with gr.Column(scale=1):
                 text_input2.render()
                 text_input3.render()
+                num_of_inf_steps_slider3.render()
                 noised_image_checkbox_2.render()
             with gr.Column(scale=1):
                 image_input4.render()
@@ -175,18 +202,20 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
                 #             examples_per_page=4)
 
     run_generate_mixed_image_button.click(run_generate_mixed_image,
-                        inputs=[text_input2,text_input3,noised_image_checkbox_2],
+                        inputs=[text_input2,text_input3,num_of_inf_steps_slider3,noised_image_checkbox_2],
                         outputs=image_input4)
 
     ####################################################################################################################
-    with gr.Tab("progress_video"):
+    with gr.Tab("Generate Image with Style"):
         # Images
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
                 text_input5.render()
+                num_of_inf_steps_slider4.render()
+                selected_style.render()
 
             with gr.Column(scale=1):
-                video_output.render()
+                image_input5.render()
 
         # Submit & Clear
         with gr.Row():
@@ -194,18 +223,18 @@ with gr.Blocks(css=css, title='Play with Stable Diffusion') as demo:
                 run_progress_video_button = gr.Button("progress_video", variant='primary')
                 clear_btn_progress_video = gr.Button("Clear", variant="secondary")
 
-                gr.Markdown(description_progress_video)
+                gr.Markdown(description_generate_prompt_with_style)
                 # gr.Examples(examples = ["examples/12830823_87d2654e31.jpg", "examples/27782020_4dab210360.jpg", "examples/44129946_9eeb385d77.jpg"],
                 #             inputs=[text_input5],
-                #             outputs=video_output,
+                #             outputs=image_input5,
                 #             fn=run_process_video,
                 #             examples_per_page=4)
 
-    run_progress_video_button.click(run_process_video,
+    run_progress_video_button.click(run_generate_image_with_style,
                         inputs=[
-                            text_input5,
+                            text_input5,selected_style,num_of_inf_steps_slider4
                         ],
-                        outputs=video_output)
+                        outputs=image_input5)
     
     #######################################################################################################################
     #######################################################################################################################

main_inference.py

@@ -10,8 +10,6 @@ from torchvision import transforms as tfms
 import shutil
 # For video display:
 import cv2
-from IPython.display import HTML
-from base64 import b64encode
 import os 
 from utils import color_loss,latents_to_pil,pil_to_latent,sketch_loss
 # Set device
@@ -34,11 +32,11 @@ unet = unet.to(torch_device)
 
 scheduler.set_timesteps(15)
 
-def generate_mixed_image(prompt1, prompt2,noised_image=False):
+def generate_mixed_image(prompt1, prompt2,num_inference_steps=50,noised_image=False):
     mix_factor = 0.4 #@param
     height = 512                        # default height of Stable Diffusion
     width = 512                         # default width of Stable Diffusion
-    num_inference_steps = 50  #@param           # Number of denoising steps
+    num_inference_steps = num_inference_steps  #@param           # Number of denoising steps
     guidance_scale = 8                # Scale for classifier-free guidance
     generator = torch.manual_seed(32)   # Seed generator to create the inital latent noise
     batch_size = 1
@@ -100,12 +98,12 @@ def generate_mixed_image(prompt1, prompt2,noised_image=False):
 
     return output
 
-def generate_image(prompt,color_postprocessing=False,postporcessing_color=None,color_loss_scale=40,noised_image=False):
+def generate_image(prompt,num_inference_steps=50,color_postprocessing=False,postporcessing_color=None,color_loss_scale=40,noised_image=False):
     #@title Store the predicted outputs and next frame for later viewing
     #prompt = 'A campfire (oil on canvas)' #
     height = 512                        # default height of Stable Diffusion
     width = 512                         # default width of Stable Diffusion
-    num_inference_steps = 50  #          # Number of denoising steps
+    num_inference_steps = num_inference_steps  #          # Number of denoising steps
     guidance_scale = 8 #         # Scale for classifier-free guidance
     generator = torch.manual_seed(32)   # Seed generator to create the inital latent noise
     batch_size = 1
@@ -210,26 +208,6 @@ def generate_image(prompt,color_postprocessing=False,postporcessing_color=None,c
 
     return output
 
-def progress_video(prompt):
-    pil_image = generate_image(prompt)
-    # Generate a list of image file paths (replace with your own logic)
-    num_frames = len(os.listdir("steps/"))
-    image_files = [f"steps/{i:04d}.jpeg" for i in range(1, num_frames + 1)]
-    # Read the first image to get its size (assuming all images have the same size)
-    first_image = cv2.imread({image_files[0]})
-    height, width, _ = first_image.shape
-
-    # Define the output video writer
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
-    out = cv2.VideoWriter('out.mp4', fourcc, 12, (width, height))
-
-    for image_file in image_files:
-        frame = cv2.imread(image_file)
-        out.write(frame)
-
-    out.release()
-    return "out.mp4"
-
 def generate_noised_version_of_image(pil_image):
     # View a noised version
     encoded = pil_to_latent(pil_image,vae)
@@ -239,6 +217,7 @@ def generate_noised_version_of_image(pil_image):
     return latents_to_pil(encoded_and_noised,vae)[0] # Display
 
 
+
 # if __name__ == "__main__":
 #     prompt = 'A campfire (oil on canvas)'
 #     color_loss_scale = 40
@@ -248,5 +227,3 @@ def generate_noised_version_of_image(pil_image):
 #     #pil_image = generate_noised_version_of_image(Image.open('output.png').resize((512, 512)))
 #     pil_image.save("output1.png")
     
-if __name__ == "__main__":
-    progress_video("lol")

requirements.txt

@@ -1,8 +1,12 @@
-transformers
-diffusers==0.2.4 
-sentence_transformers
-gradio
-torch
-torchvision
-matplotlib
-opencv-python
+diffusers==0.21.4
+gradio==3.49.0
+ipython==8.15.0
+matplotlib==3.8.0
+numpy==1.26.1
+opencv_python==4.8.1.78
+Pillow==10.1.0
+Pillow==10.1.0
+torch==2.1.0
+torchvision==0.16.0
+tqdm==4.66.1
+transformers==4.25.1

style_embeddings/learned_embeds_arcane.bin

@@ -0,0 +1 @@
+{"payload":{"allShortcutsEnabled":false,"fileTree":{"S20":{"items":[{"name":"README.md","path":"S20/README.md","contentType":"file"},{"name":"S20-HuggingFace.ipynb","path":"S20/S20-HuggingFace.ipynb","contentType":"file"},{"name":"S20.ipynb","path":"S20/S20.ipynb","contentType":"file"},{"name":"learned_embeds_arcane.bin","path":"S20/learned_embeds_arcane.bin","contentType":"file"},{"name":"learned_embeds_dr_strange.bin","path":"S20/learned_embeds_dr_strange.bin","contentType":"file"},{"name":"learned_embeds_gta5.bin","path":"S20/learned_embeds_gta5.bin","contentType":"file"},{"name":"learned_embeds_illustration.bin","path":"S20/learned_embeds_illustration.bin","contentType":"file"},{"name":"learned_embeds_pokemon.bin","path":"S20/learned_embeds_pokemon.bin","contentType":"file"}],"totalCount":8},"":{"items":[{"name":"S10","path":"S10","contentType":"directory"},{"name":"S11","path":"S11","contentType":"directory"},{"name":"S12","path":"S12","contentType":"directory"},{"name":"S13","path":"S13","contentType":"directory"},{"name":"S15","path":"S15","contentType":"directory"},{"name":"S16","path":"S16","contentType":"directory"},{"name":"S18","path":"S18","contentType":"directory"},{"name":"S19","path":"S19","contentType":"directory"},{"name":"S20","path":"S20","contentType":"directory"},{"name":"S5","path":"S5","contentType":"directory"},{"name":"S6","path":"S6","contentType":"directory"},{"name":"S7","path":"S7","contentType":"directory"},{"name":"S8","path":"S8","contentType":"directory"},{"name":"S9","path":"S9","contentType":"directory"},{"name":"README.md","path":"README.md","contentType":"file"}],"totalCount":15}},"fileTreeProcessingTime":5.458002,"foldersToFetch":[],"reducedMotionEnabled":null,"repo":{"id":648622518,"defaultBranch":"main","name":"ERA-V1-Assignments","ownerLogin":"gupta1912","currentUserCanPush":false,"isFork":false,"isEmpty":false,"createdAt":"2023-06-02T12:02:53.000Z","ownerAvatar":"https://avatars.githubusercontent.com/u/47393639?v=4","public":true,"private":false,"isOrgOwned":false},"symbolsExpanded":false,"treeExpanded":true,"refInfo":{"name":"main","listCacheKey":"v0:1685707373.878437","canEdit":false,"refType":"branch","currentOid":"69279d066a624710a5d9143b24292e736177714d"},"path":"S20/learned_embeds_arcane.bin","currentUser":null,"blob":{"rawLines":null,"stylingDirectives":null,"csv":null,"csvError":null,"dependabotInfo":{"showConfigurationBanner":false,"configFilePath":null,"networkDependabotPath":"/gupta1912/ERA-V1-Assignments/network/updates","dismissConfigurationNoticePath":"/settings/dismiss-notice/dependabot_configuration_notice","configurationNoticeDismissed":null,"repoAlertsPath":"/gupta1912/ERA-V1-Assignments/security/dependabot","repoSecurityAndAnalysisPath":"/gupta1912/ERA-V1-Assignments/settings/security_analysis","repoOwnerIsOrg":false,"currentUserCanAdminRepo":false},"displayName":"learned_embeds_arcane.bin","displayUrl":"https://github.com/gupta1912/ERA-V1-Assignments/blob/main/S20/learned_embeds_arcane.bin?raw=true","headerInfo":{"blobSize":"3.73 KB","deleteInfo":{"deleteTooltip":"You must be signed in to make or propose changes"},"editInfo":{"editTooltip":"You must be signed in to make or propose changes"},"ghDesktopPath":"https://desktop.github.com","gitLfsPath":null,"onBranch":true,"shortPath":"3672956","siteNavLoginPath":"/login?return_to=https%3A%2F%2Fgithub.com%2Fgupta1912%2FERA-V1-Assignments%2Fblob%2Fmain%2FS20%2Flearned_embeds_arcane.bin","isCSV":false,"isRichtext":false,"toc":null,"lineInfo":{"truncatedLoc":null,"truncatedSloc":null},"mode":"file"},"image":false,"isCodeownersFile":null,"isPlain":false,"isValidLegacyIssueTemplate":false,"issueTemplateHelpUrl":"https://docs.github.com/articles/about-issue-and-pull-request-templates","issueTemplate":null,"discussionTemplate":null,"language":null,"languageID":null,"large":false,"loggedIn":false,"newDiscussionPath":"/gupta1912/ERA-V1-Assignments/discussions/new","newIssuePath":"/gupta1912/ERA-V1-Assignments/issues/new","planSupportInfo":{"repoIsFork":null,"repoOwnedByCurrentUser":null,"requestFullPath":"/gupta1912/ERA-V1-Assignments/blob/main/S20/learned_embeds_arcane.bin","showFreeOrgGatedFeatureMessage":null,"showPlanSupportBanner":null,"upgradeDataAttributes":null,"upgradePath":null},"publishBannersInfo":{"dismissActionNoticePath":"/settings/dismiss-notice/publish_action_from_dockerfile","dismissStackNoticePath":"/settings/dismiss-notice/publish_stack_from_file","releasePath":"/gupta1912/ERA-V1-Assignments/releases/new?marketplace=true","showPublishActionBanner":false,"showPublishStackBanner":false},"rawBlobUrl":"https://github.com/gupta1912/ERA-V1-Assignments/raw/main/S20/learned_embeds_arcane.bin","renderImageOrRaw":true,"richText":null,"renderedFileInfo":null,"shortPath":null,"tabSize":8,"topBannersInfo":{"overridingGlobalFundingFile":false,"globalPreferredFundingPath":null,"repoOwner":"gupta1912","repoName":"ERA-V1-Assignments","showInvalidCitationWarning":false,"citationHelpUrl":"https://docs.github.com/en/github/creating-cloning-and-archiving-repositories/creating-a-repository-on-github/about-citation-files","showDependabotConfigurationBanner":false,"actionsOnboardingTip":null},"truncated":false,"viewable":false,"workflowRedirectUrl":null,"symbols":null},"copilotInfo":null,"csrf_tokens":{"/gupta1912/ERA-V1-Assignments/branches":{"post":"BDalnmGaFsCat0pyEB2XRtEEgU1tU3hPkqod54YEtINRMIyVpZhY2QaEfjaGTzB49WIGe8mvKtoNNi3HfJg_fA"},"/repos/preferences":{"post":"9b2BsWn9aEPKIJwMKJX1lraH9ktRNImj8HovOQHsPkUk6M41p49XqC_qiNv0dz03T8TyuJBk3-bJqf53V9cV9g"}}},"title":"ERA-V1-Assignments/S20/learned_embeds_arcane.bin at main · gupta1912/ERA-V1-Assignments"}

style_guidence.py

@@ -0,0 +1,176 @@
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
+from tqdm.auto import tqdm
+from torch import autocast
+from PIL import Image
+from matplotlib import pyplot as plt
+import numpy
+from torchvision import transforms as tfms
+import shutil
+# For video display:
+import cv2
+from IPython.display import HTML
+from base64 import b64encode
+import os 
+from utils import color_loss,pil_to_latent,sketch_loss
+# Set device
+torch_device =  "cpu"
+
+vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
+
+# Load the tokenizer and text encoder to tokenize and encode the text.
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+# The UNet model for generating the latents.
+unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
+
+# The noise scheduler
+scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+vae = vae.to(torch_device)
+text_encoder = text_encoder.to(torch_device)
+unet = unet.to(torch_device)
+
+scheduler.set_timesteps(15)
+
+token_emb_layer = text_encoder.text_model.embeddings.token_embedding
+pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
+
+position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
+position_embeddings = pos_emb_layer(position_ids)
+     
+     
+
+def get_output_embeds(input_embeddings):
+    # CLIP's text model uses causal mask, so we prepare it here:
+    bsz, seq_len = input_embeddings.shape[:2]
+    causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)
+
+    # Getting the output embeddings involves calling the model with passing output_hidden_states=True
+    # so that it doesn't just return the pooled final predictions:
+    encoder_outputs = text_encoder.text_model.encoder(
+        inputs_embeds=input_embeddings,
+        attention_mask=None, # We aren't using an attention mask so that can be None
+        causal_attention_mask=causal_attention_mask.to(torch_device),
+        output_attentions=None,
+        output_hidden_states=True, # We want the output embs not the final output
+        return_dict=None,
+    )
+
+    # We're interested in the output hidden state only
+    output = encoder_outputs[0]
+
+    # There is a final layer norm we need to pass these through
+    output = text_encoder.text_model.final_layer_norm(output)
+
+    # And now they're ready!
+    return output
+
+def set_timesteps(scheduler, num_inference_steps):
+    scheduler.set_timesteps(num_inference_steps)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+# def latents_to_pil(latents):
+#     # bath of latents -> list of images
+#     latents = (1 / 0.18215) * latents
+#     with torch.no_grad():
+#         image = vae.decode(latents).sample
+#     image = (image / 2 + 0.5).clamp(0, 1)
+#     image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+#     images = (image * 255).round().astype("uint8")
+#     pil_images = [Image.fromarray(image) for image in images]
+#     return pil_images
+
+def latents_to_pil(latents):
+      # bath of latents -> list of images
+  latents = (1 / 0.18215) * latents
+  with torch.no_grad():
+    image = vae.decode(latents)
+  image = (image / 2 + 0.5).clamp(0, 1)
+  image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+  images = (image * 255).round().astype("uint8")
+  pil_images = [Image.fromarray(image) for image in images]
+  return pil_images
+
+
+def generate_with_embs(text_embeddings,text_input, seed,num_inference_steps):
+    
+    height = 512                        # default height of Stable Diffusion
+    width = 512                         # default width of Stable Diffusion
+    num_inference_steps = num_inference_steps        # Number of denoising steps
+    guidance_scale = 7.5                # Scale for classifier-free guidance
+    generator = torch.manual_seed(seed)   # Seed generator to create the inital latent noise
+    batch_size = 1
+
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = tokenizer(
+      [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
+    )
+    with torch.no_grad():
+        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
+    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+    # Prep Scheduler
+    set_timesteps(scheduler, num_inference_steps)
+
+    # Prep latents
+    latents = torch.randn(
+    (batch_size, unet.in_channels, height // 8, width // 8),
+    generator=generator,
+    )
+    latents = latents.to(torch_device)
+    # latents = latents * scheduler.init_noise_sigma
+    latents = latents * scheduler.sigmas[0] # Need to scale to match k
+
+    # Loop
+    for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latent_model_input = torch.cat([latents] * 2)
+        sigma = scheduler.sigmas[i]
+        #latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+        latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
+        # predict the noise residual
+        with torch.no_grad():
+            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        # compute the previous noisy sample x_t -> x_t-1
+        #latents = scheduler.step(noise_pred, t, latents).prev_sample
+        latents = scheduler.step(noise_pred, i, latents)["prev_sample"]
+    return latents_to_pil(latents)[0]
+
+def generate_with_prompt_style(prompt, style, num_of_inf_steps=50,seed = 42):
+
+    prompt = prompt + ' in style of s'
+    embed = torch.load(style)
+    print("Keys",embed.keys())
+    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    # for t in text_input['input_ids'][0][:20]: # We'll just look at the first 7 to save you from a wall of '<|endoftext|>'
+    #     print(t, tokenizer.decoder.get(int(t)))
+    input_ids = text_input.input_ids.to(torch_device)
+
+    token_embeddings = token_emb_layer(input_ids)
+    # The new embedding - our special birb word
+    replacement_token_embedding = embed[list(embed.keys())[0]].to(torch_device)
+
+    # Insert this into the token embeddings
+    token_embeddings[0, torch.where(input_ids[0]==338)] = replacement_token_embedding.to(torch_device)
+
+    # Combine with pos embs
+    input_embeddings = token_embeddings + position_embeddings
+
+    #  Feed through to get final output embs
+    modified_output_embeddings = get_output_embeds(input_embeddings)
+
+    # And generate an image with this:
+    return generate_with_embs(modified_output_embeddings, text_input, seed,num_of_inf_steps)
+
+
+# prompt = 'A man sipping wine wearing a spacesuit on the moon'
+# image = generate_with_prompt_style(prompt, '/home/deepanshudashora/Documents/Stable_Diffusion/caitlin_fairchild.bin')
+
+# image.save("output.png")