Update app.py

app.py

@@ -11,12 +11,18 @@ import timm
 # model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
 # model.train()
 
-model = BadNet(3, 10)
+# model = BadNet(3, 10)
 
+from diffusers import DiffusionPipeline
+
+pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+# pipeline = pipeline.to('cuda:0')
 
 
 import os 
 
+
+
 def print_bn():
     bn_data = []
     for m in model.modules():
@@ -27,7 +33,15 @@ def print_bn():
             bn_data.append(m.momentum)
     return bn_data
 
-def greet(image):
+def greet(text):
+    if(text is None):
+        pipeline.unet.load_attn_procs('./models/pytorch_lora_weights.bin')
+    else:
+        images = pipeline(text).images
+        image = images[0]
+        return image
+
+ def greet_old(image):
     # url = f'https://huggingface.co/spaces?p=1&sort=modified&search=GPT'
     # html = request_url(url)
     # key = os.getenv("OPENAI_API_KEY")
@@ -68,6 +82,9 @@ def greet(image):
     return out
 
 
-image = gr.inputs.Image(label="Upload a photo", shape=(32,32))
-iface = gr.Interface(fn=greet, inputs=image, outputs="text")
+
+iface = gr.Interface(fn=greet, inputs='text', output="image")
+
+# image = gr.inputs.Image(label="Upload a photo", shape=(32,32))
+# iface = gr.Interface(fn=greet, inputs=image, outputs="text")
 iface.launch()