Update app.py

app.py

@@ -5,182 +5,158 @@ import requests
 import uuid
 import io 
 import base64
-import random
-from transforms import RGBTransform # from source code mentioned above
+from transforms import RGBTransform
+import concurrent.futures
+import time
 
+# Dictionary to track model availability status
+model_status = {}
 
-loaded_model=[]
-for i,model in enumerate(models):
-    try:
-        loaded_model.append(gr.load(f'models/{model}'))
-    except Exception as e:
-        print(e)
-        pass
-print (loaded_model)
-
-
-
-def run_dif_color(out_prompt,model_drop,cnt,color,tint):
-    h=color.lstrip('#')
-    #h = input('Enter hex: ').lstrip('#')
-    #print('RGB =', tuple(int(h[i:i+2], 16) for i in (0, 2, 4)))
-    color=tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
-    
-    print (color)
-    p_seed=""
-    out_box=[]
-    out_html=""
-    #for i,ea in enumerate(loaded_model):
-    
-    
-    
-    for i in range(int(cnt)):
-        rand=random.randint(1,500)
-        for i in range(rand):
-            p_seed+=" "
-        try:
-            #model=gr.load(f'models/{model[int(model_drop)]}')        
-            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(color,factor=float(tint)).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=[]
-    out_html=""
-    #for i,ea in enumerate(loaded_model):
-    for i in range(int(cnt)):
-        p_seed+=" "
+def load_models():
+    """
+    Attempts to load all models and tracks their availability status
+    Returns a list of successfully loaded models
+    """
+    loaded_models = []
+    for model in models:
         try:
-            model=loaded_model[int(model_drop)]
-            out_img=model(out_prompt+p_seed)
-            print(out_img)
-            out_box.append(out_img)           
+            # Attempt to load the model
+            loaded_model = gr.load(f'models/{model}')
+            loaded_models.append(loaded_model)
+            model_status[model] = {'status': 'available', 'error': None}
         except Exception as e:
-            print(e)
-            out_html=str(e)
-            pass
-        yield out_box,out_html
-
-def run_dif_og(out_prompt,model_drop,cnt):
-    out_box=[]
-    out_html=""
-    #for i,ea in enumerate(loaded_model):
-    for i in range(cnt):
-        try:
-            #print (ea)
-            model=loaded_model[int(model_drop)]
-            out_img=model(out_prompt)
-            print(out_img)
-            url=f'https://omnibus-top-20.hf.space/file={out_img}'
-            print(url)
-            uid = uuid.uuid4()
-            #urllib.request.urlretrieve(image, 'tmp.png')
-            #out=Image.open('tmp.png')
-            r = requests.get(url, stream=True)
-            
-            if r.status_code == 200:
-                img_buffer = io.BytesIO(r.content)
-                print (f'bytes:: {io.BytesIO(r.content)}')
-                str_equivalent_image = base64.b64encode(img_buffer.getvalue()).decode()
-                img_tag = "<img src='data:image/png;base64," + str_equivalent_image + "'/>"                
-                out_html+=f"<div  class='img_class'><a href='https://huggingface.co/models/{models[i]}'>{models[i]}</a><br>"+img_tag+"</div>"
-                out = Image.open(io.BytesIO(r.content))
-                out_box.append(out)
-            html_out = "<div class='grid_class'>"+out_html+"</div>"
-            yield out_box,html_out
-        except Exception as e:
-            out_html+=str(e)
-            html_out = "<div class='grid_class'>"+out_html+"</div>"
-            
-            yield out_box,html_out
+            # Track failed model loads
+            model_status[model] = {'status': 'unavailable', 'error': str(e)}
+            print(f"Failed to load {model}: {e}")
+    return loaded_models
 
-def thread_dif(out_prompt,mod):
-    out_box=[]
-    out_html=""
-    #for i,ea in enumerate(loaded_model):
+def generate_single_image(model_name, model, prompt, color=None, tint_strength=0.3):
+    """
+    Generates a single image from a specific model with optional color tinting
+    Returns tuple of (image, error_message, model_name)
+    """
     try:
-        print (ea)
-        model=loaded_model[int(mod)]
-        out_img=model(out_prompt)
-        print(out_img)
-        url=f'https://omnibus-top-20.hf.space/file={out_img}'
-        print(url)
-        uid = uuid.uuid4()
-        #urllib.request.urlretrieve(image, 'tmp.png')
-        #out=Image.open('tmp.png')
-        r = requests.get(url, stream=True)
+        # Generate image
+        out_img = model(prompt)
         
-        if r.status_code == 200:
-            img_buffer = io.BytesIO(r.content)
-            print (f'bytes:: {io.BytesIO(r.content)}')
-            str_equivalent_image = base64.b64encode(img_buffer.getvalue()).decode()
-            img_tag = "<img src='data:image/png;base64," + str_equivalent_image + "'/>"                
-            
-            #out_html+=f"<div  class='img_class'><a href='https://huggingface.co/models/{models[i]}'>{models[i]}</a><br>"+img_tag+"</div>"
-            out = Image.open(io.BytesIO(r.content))
-            out_box.append(out)
+        # Process the image
+        if isinstance(out_img, str):  # If URL is returned
+            r = requests.get(f'https://omnibus-top-20.hf.space/file={out_img}', stream=True)
+            if r.status_code != 200:
+                return None, f"HTTP Error: {r.status_code}", model_name
+                
+            img = Image.open(io.BytesIO(r.content)).convert('RGB')
         else:
-            out_html=r.status_code
-        html_out = "<div class='grid_class'>"+out_html+"</div>"
-        return out_box,html_out
-    except Exception as e:
-        out_html=str(e)
-        #out_html+=str(e)
-        html_out = "<div class='grid_class'>"+out_html+"</div>"
+            img = Image.open(out_img).convert('RGB')
+            
+        # Apply color tinting if specified
+        if color is not None:
+            h = color.lstrip('#')
+            rgb_color = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
+            img = RGBTransform().mix_with(rgb_color, factor=float(tint_strength)).applied_to(img)
+            
+        return img, None, model_name
         
-        return out_box,html_out
+    except Exception as e:
+        return None, str(e), model_name
 
+def run_all_models(prompt, color=None, tint_strength=0.3):
+    """
+    Generates images from all available models in parallel
+    """
+    results = []
+    errors = []
+    
+    # Load models if not already loaded
+    loaded_models = load_models()
+    
+    # Use ThreadPoolExecutor for parallel execution
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+        future_to_model = {
+            executor.submit(
+                generate_single_image, 
+                model_name, 
+                model, 
+                prompt, 
+                color, 
+                tint_strength
+            ): model_name 
+            for model_name, model in zip(models, loaded_models)
+        }
+        
+        for future in concurrent.futures.as_completed(future_to_model):
+            img, error, model_name = future.result()
+            if error:
+                errors.append(f"{model_name}: {error}")
+                model_status[model_name]['status'] = 'failed'
+                model_status[model_name]['error'] = error
+            if img:
+                results.append((img, model_name))
+                
+    # Generate HTML report
+    html_report = "<div class='results-grid'>"
+    for model in models:
+        status = model_status[model]
+        status_color = {
+            'available': 'green',
+            'unavailable': 'red',
+            'failed': 'orange'
+        }.get(status['status'], 'gray')
+        
+        html_report += f"""
+        <div class='model-status'>
+            <h3>{model}</h3>
+            <p style='color: {status_color}'>Status: {status['status']}</p>
+            {f"<p class='error'>Error: {status['error']}</p>" if status['error'] else ""}
+        </div>
+        """
+    html_report += "</div>"
+    
+    return results, html_report
 
-            
-def start_threads(prompt):
-    t1 = threading.Thread(target=thread_dif, args=(prompt,0))
-    t2 = threading.Thread(target=thread_dif, args=(prompt,1))
-    t1.start()
-    t2.start()
-    print (t1)
-    print (t2)
-    a1,a2=t1.result()
-    b1,b2=t2.result()
-    return a1,a2
-            
-css="""
-.grid_class{
-display:flex;
-height:100%;
+# Gradio interface
+css = """
+.results-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
+    gap: 1rem;
+    padding: 1rem;
 }
-.img_class{
-min-width:200px;
+.model-status {
+    border: 1px solid #ddd;
+    padding: 1rem;
+    border-radius: 4px;
+}
+.error {
+    color: red;
+    font-size: 0.9em;
+    word-break: break-word;
 }
-
 """
+
 with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as app:
     with gr.Row():
         with gr.Column():
-            inp=gr.Textbox(label="Prompt")
-            btn=gr.Button()
+            inp = gr.Textbox(label="Prompt")
+            btn = gr.Button("Generate from All Models")
         with gr.Column():
-            col = gr.ColorPicker(label="Color Tint")
+            col = gr.ColorPicker(label="Color Tint (Optional)")
             tint = gr.Slider(label="Tint Strength", minimum=0, maximum=1, step=0.01, value=0.30)
-        
-    with gr.Row():
-        model_drop=gr.Dropdown(label="Models", choices=models, type='index', value=models[0])
-        cnt = gr.Number(value=1)
-    out_html=gr.HTML()
-    outp=gr.Gallery()
-    btn.click(run_dif_color,[inp,model_drop,cnt,col,tint],[outp,out_html])
+    
+    status_html = gr.HTML(label="Model Status")
+    gallery = gr.Gallery()
+    
+    def process_and_display(prompt, color, tint_strength):
+        results, html_report = run_all_models(prompt, color, tint_strength)
+        return (
+            [img for img, _ in results],
+            html_report
+        )
+    
+    btn.click(
+        process_and_display,
+        inputs=[inp, col, tint],
+        outputs=[gallery, status_html]
+    )
+
 app.launch()