Update app.py

app.py

@@ -5,7 +5,8 @@ import requests
 import uuid
 import io 
 import base64
-import threading
+from transforms import RGBTransform # from source code mentioned above
+
 
 loaded_model=[]
 for i,model in enumerate(models):
@@ -16,6 +17,35 @@ for i,model in enumerate(models):
         pass
 print (loaded_model)
 
+
+
+def run_dif_color(out_prompt,model_drop,cnt,color):
+
+    p_seed=""
+    out_box=[]
+    out_html=""
+    #for i,ea in enumerate(loaded_model):
+    for i in range(int(cnt)):
+        p_seed+=" "
+        try:
+            model=loaded_model[int(model_drop)]
+            out_img=model(out_prompt+p_seed)
+            print(out_img)
+
+            raw=Image.open(out_img)
+            raw=raw.convert('RGB')
+
+            colorize = RGBTransform().mix_with((255, 0, 0),factor=.30).applied_to(raw)    
+            
+            out_box.append(colorize)           
+        except Exception as e:
+            print(e)
+            out_html=str(e)
+            pass
+            
+        yield out_box,out_html
+
+
 def run_dif(out_prompt,model_drop,cnt):
     p_seed=""
     out_box=[]
@@ -136,5 +166,5 @@ with gr.Blocks(css=css) as app:
         cnt = gr.Number(value=1)
     out_html=gr.HTML()
     outp=gr.Gallery()
-    btn.click(run_dif,[inp,model_drop,cnt],[outp,out_html])
+    btn.click(run_dif_color,[inp,model_drop,cnt],[outp,out_html])
 app.launch()