Update app.py

app.py

@@ -1,162 +1,127 @@
 import gradio as gr
-from models import models
 from PIL import Image
 import requests
-import uuid
-import io 
+import io
 import base64
-from transforms import RGBTransform
-import concurrent.futures
-import time
 
-# Dictionary to track model availability status
-model_status = {}
+# List of models to process
+models = [
+    "dreamlike-art/dreamlike-photoreal-2.0",
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "black-forest-labs/FLUX.1-dev",
+    "black-forest-labs/FLUX.1-schnell",
+    "veryVANYA/ps1-style-flux",
+    "alvdansen/softserve_anime",
+    "multimodalart/flux-tarot-v1",
+    "davisbro/half_illustration",
+    "dataautogpt3/OpenDalleV1.1",
+    "aleksa-codes/flux-ghibsky-illustration",
+    "alvdansen/flux-koda",
+    "openskyml/soviet-diffusion-xl",
+    "XLabs-AI/flux-RealismLora",
+    "alvdansen/frosting_lane_flux",
+    "alvdansen/phantasma-anime",
+    "kudzueye/Boreal",
+    "glif/how2draw",
+    "dataautogpt3/FLUX-AestheticAnime",
+    "prithivMLmods/Fashion-Hut-Modeling-LoRA",
+    "dataautogpt3/FLUX-SyntheticAnime",
+    "brushpenbob/flux-midjourney-anime",
+    "robert123231/coloringbookgenerator",
+    "prithivMLmods/Castor-Collage-Dim-Flux-LoRA",
+    "prithivMLmods/Flux-Product-Ad-Backdrop",
+    "multimodalart/product-design",
+    "glif/90s-anime-art",
+    "glif/Brain-Melt-Acid-Art",
+    "lustlyai/Flux_Lustly.ai_Uncensored_nsfw_v1",
+    "Keltezaa/NSFW_MASTER_FLUX",
+    "tryonlabs/FLUX.1-dev-LoRA-Outfit-Generator",
+    "Jovie/Midjourney",
+    "Yntec/DreamPhotoGASM",
+    "strangerzonehf/Flux-Super-Realism-LoRA",
+    "stabilityai/stable-diffusion-2-1-base",
+    "stabilityai/stable-diffusion-3.5-large",
+    "stabilityai/stable-diffusion-3.5-large-turbo",
+    "stabilityai/stable-diffusion-3-medium-diffusers",
+    "stablediffusionapi/duchaiten-real3d-nsfw-xl",
+    "nerijs/pixel-art-xl",
+    "KappaNeuro/character-design",
+    "alvdansen/sketchedoutmanga",
+    "alvdansen/archfey_anime",
+    "alvdansen/lofi-cuties",
+    "Yntec/YiffyMix",
+    "digiplay/AnalogMadness-realistic-model-v7",
+    "artificialguybr/selfiephotographyredmond-selfie-photography-lora-for-sdxl",
+    "artificialguybr/filmgrain-redmond-filmgrain-lora-for-sdxl",
+    "goofyai/Leonardo_Ai_Style_Illustration",
+    "goofyai/cyborg_style_xl",
+    "alvdansen/littletinies",
+    "Dremmar/nsfw-xl",
+    "artificialguybr/analogredmond",
+    "artificialguybr/PixelArtRedmond",
+    "CiroN2022/ascii-art",
+    "Yntec/Analog",
+    "Yntec/MapleSyrup",
+    "digiplay/perfectLewdFantasy_v1.01",
+    "digiplay/AbsoluteReality_v1.8.1",
+    "goofyai/disney_style_xl",
+    "artificialguybr/LogoRedmond-LogoLoraForSDXL-V2",
+    "Yntec/epiCPhotoGasm",
+]
 
-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:
+# Initialize a dictionary to track models and their status
+model_results = {"success": {}, "failed": []}
+
+def generate_images(prompt):
+    output_images = []
+    failed_models = []
+    for model_name in models:
         try:
-            # Attempt to load the model
-            loaded_model = gr.load(f'models/{model}')
-            loaded_models.append(loaded_model)
-            model_status[model] = {'status': 'available', 'error': None}
+            # Attempt to load and generate an image with the model
+            model = gr.Interface.load(f"models/{model_name}")
+            img_path = model.predict(prompt)
+
+            # Fetch image
+            response = requests.get(img_path, stream=True)
+            if response.status_code == 200:
+                img = Image.open(io.BytesIO(response.content))
+                output_images.append((model_name, img))
+            else:
+                failed_models.append(model_name)
         except Exception as e:
-            # Track failed model loads
-            model_status[model] = {'status': 'unavailable', 'error': str(e)}
-            print(f"Failed to load {model}: {e}")
-    return loaded_models
+            print(f"Error with model {model_name}: {e}")
+            failed_models.append(model_name)
 
-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:
-        # Generate image
-        out_img = model(prompt)
-        
-        # 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:
-            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
-        
-    except Exception as e:
-        return None, str(e), model_name
+    return output_images, failed_models
 
-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 ""}
+def app_interface(prompt):
+    output_images, failed_models = generate_images(prompt)
+    images_html = ""
+    for model_name, img in output_images:
+        # Convert images for display in Gradio
+        img_buffer = io.BytesIO()
+        img.save(img_buffer, format="PNG")
+        img_str = base64.b64encode(img_buffer.getvalue()).decode()
+        images_html += f"""
+        <div>
+            <strong>{model_name}</strong>
+            <img src='data:image/png;base64,{img_str}' style="max-width: 200px;" />
         </div>
         """
-    html_report += "</div>"
     
-    return results, html_report
-
-# Gradio interface
-css = """
-.results-grid {
-    display: grid;
-    grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
-    gap: 1rem;
-    padding: 1rem;
-}
-.model-status {
-    border: 1px solid #ddd;
-    padding: 1rem;
-    border-radius: 4px;
-}
-.error {
-    color: red;
-    font-size: 0.9em;
-    word-break: break-word;
-}
-"""
+    failed_html = "<ul>" + "".join(f"<li>{model}</li>" for model in failed_models) + "</ul>"
+    return images_html, f"<h3>Failed Models:</h3>{failed_html}"
 
-with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as app:
+# Define Gradio UI
+with gr.Blocks() as app:
     with gr.Row():
         with gr.Column():
-            inp = gr.Textbox(label="Prompt")
-            btn = gr.Button("Generate from All Models")
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt")
+            generate_button = gr.Button("Generate Images")
         with gr.Column():
-            col = gr.ColorPicker(label="Color Tint (Optional)")
-            tint = gr.Slider(label="Tint Strength", minimum=0, maximum=1, step=0.01, value=0.30)
-    
-    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]
-    )
+            output_html = gr.HTML(label="Generated Images")
+            failed_html = gr.HTML(label="Failed Models")
+
+    generate_button.click(app_interface, inputs=[prompt_input], outputs=[output_html, failed_html])
 
 app.launch()