app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from io import BytesIO
 from random import randint
 from all_models import models
 
@@ -28,8 +29,21 @@ def update_imgbox(choices):
 def gen_fn(model_str, prompt):
     if model_str == 'NA':
         return None
-    noise = str(randint(0, 9999999))
-    return models_load[model_str](f'{prompt} {noise}')
+    noise = str(randint(0, 99999999))
+    result = models_load[model_str](f'{prompt} {noise}')
+    
+    if isinstance(result, str):
+        response = requests.get(result)
+        img = Image.open(BytesIO(response.content))
+    elif isinstance(result, bytes):
+        img = Image.open(BytesIO(result))
+    else:
+        return None
+    
+    buffered = BytesIO()
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return f"data:image/png;base64,{img_str}"
 
 def make_me():
     with gr.Row():
@@ -152,9 +166,11 @@ body {
     border: 1px solid #3b4252;
     background-color: #2d343f;
     border-radius: 4px;
-    margin: 10px;
+    margin: 20px;
     max-width: 100%;
     box-sizing: border-box;
+    width: 100%;
+    height: auto;
 }
 .custom_accordion {
     background-color: #2d3d4f;
@@ -217,5 +233,5 @@ body {
 with gr.Blocks(css=custom_css) as demo: 
     make_me()
 
-demo.queue(concurrency_count=10)
+demo.queue(concurrency_count=50)
 demo.launch()