Update app.py

app.py

@@ -8,12 +8,17 @@ import gradio
 import torch
 from diffusers import StableDiffusionPipeline
 from torch import autocast
+#from PIL import Image
+#from torchvision import transforms
+
+#from diffusers import StableDiffusionImageVariationPipeline
 
 
 openai.api_key = os.getenv('openaikey')
 
 def predict(input, manual_query_repacement, history=[]):
 
+    # gpt3
     if manual_query_repacement != "":
         input = manual_query_repacement
     
@@ -30,15 +35,64 @@ def predict(input, manual_query_repacement, history=[]):
     responseText = response["choices"][0]["text"]
     history.append((input, responseText))
 
-    return history, history
+
+    #img generation
+    prompt = "Yoda"
+    scale = 10
+    n_samples = 4
+
+    # Sometimes the nsfw checker is confused by the Naruto images, you can disable
+    # it at your own risk here
+    #disable_safety = False
+
+    #if disable_safety:
+    #  def null_safety(images, **kwargs):
+    #      return images, False
+    #  pipe.safety_checker = null_safety
+
+    with autocast("cuda"):
+      images = pipe(n_samples*[prompt], guidance_scale=scale).images
+
+    for idx, im in enumerate(images):
+      im.save(f"{idx:06}.png")
+
+    images_list = pipe(
+        inp.tile(n_samples, 1, 1, 1),
+        guidance_scale=scale,
+        num_inference_steps=steps,
+        generator=generator,
+        )
+
+    images = []
+    for i, image in enumerate(images_list["images"]):
+        if(images_list["nsfw_content_detected"][i]):
+            safe_image = Image.open(r"unsafe.png")
+            images.append(safe_image)
+        else:
+            images.append(image)
+
+
+    
+    return history, history, images
 
 
 
 inputText = gradio.Textbox(value="tmp")
 manual_query = gradio.Textbox(placeholder="Input any query here, to replace the image generation query builder entirely.")
 
+output_img = gr.Gallery(label="Generated image")
+output_img.style(grid=2)
+
 gradio.Interface(fn=predict,
              inputs=[inputText,manual_query,'state'],
             
-             outputs=["chatbot",'state']).launch()
+             outputs=["chatbot",'state', output_img]).launch()
+
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = StableDiffusionPipeline.from_pretrained("stale2000/sd-dnditem", torch_dtype=torch.float16)  
+pipe = pipe.to(device)
+
+