Update app.py

app.py

@@ -1,89 +1,13 @@
-import logging
-import random
-import warnings
-import os
-import gradio as gr
-import numpy as np
-import spaces
-import torch
-from gradio_imageslider import ImageSlider
-from PIL import Image
-from huggingface_hub import hf_hub_download
-import subprocess
-import sys
-import tempfile
-from typing import Sequence, Mapping, Any, Union
-import asyncio
-import execution
-from nodes import init_extra_nodes
-import server
-
-# Copy functions from FluxSimpleUpscaler.txt
-def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    try:
-        return obj[index]
-    except KeyError:
-        return obj["result"][index]
-
-def find_path(name: str, path: str = None) -> str:
-    if path is None:
-        path = os.getcwd()
-    if name in os.listdir(path):
-        path_name = os.path.join(path, name)
-        print(f"{name} found: {path_name}")
-        return path_name
-    parent_directory = os.path.dirname(path)
-    if parent_directory == path:
-        return None
-    return find_path(name, parent_directory)
-
-def add_comfyui_directory_to_sys_path() -> None:
-    comfyui_path = find_path("ComfyUI")
-    if comfyui_path is not None and os.path.isdir(comfyui_path):
-        sys.path.append(comfyui_path)
-        print(f"'{comfyui_path}' added to sys.path")
-
-def add_extra_model_paths() -> None:
-    try:
-        from main import load_extra_path_config
-    except ImportError:
-        print("Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead.")
-        from utils.extra_config import load_extra_path_config
-    extra_model_paths = find_path("extra_model_paths.yaml")
-    if extra_model_paths is not None:
-        load_extra_path_config(extra_model_paths)
-    else:
-        print("Could not find the extra_model_paths config file.")
-
-def import_custom_nodes() -> None:
-    import asyncio
-    import execution
-    from nodes import init_extra_nodes
-    import server
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    server_instance = server.PromptServer(loop)
-    execution.PromptQueue(server_instance)
-    init_extra_nodes()
-
-# Setup ComfyUI and custom nodes
-if not os.path.exists("ComfyUI"):
-    subprocess.run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI.git"])
-
-custom_node_path = "ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale"
-if not os.path.exists(custom_node_path):
-    subprocess.run(["git", "clone", "https://github.com/ssitu/ComfyUI_UltimateSDUpscale.git", custom_node_path])
-
 # Create model directories
-os.makedirs("ComfyUI/models/unet", exist_ok=True)
+os.makedirs("ComfyUI/models/diffusion_models", exist_ok=True)
 os.makedirs("ComfyUI/models/clip", exist_ok=True)
 os.makedirs("ComfyUI/models/vae", exist_ok=True)
 os.makedirs("ComfyUI/models/upscale_models", exist_ok=True)
 
 # Download models if not present
-unet_path = "ComfyUI/models/unet/flux1-dev-fp8.safetensors"
-if not os.path.exists(unet_path):
-    hf_hub_download("Kijai/flux-fp8", "flux1-dev-fp8.safetensors", local_dir="ComfyUI/models/unet")
+diffusion_path = "ComfyUI/models/diffusion_models/flux1-dev-fp8.safetensors"
+if not os.path.exists(diffusion_path):
+    hf_hub_download("Kijai/flux-fp8", "flux1-dev-fp8.safetensors", local_dir="ComfyUI/models/diffusion_models")
 
 clip_l_path = "ComfyUI/models/clip/clip_l.safetensors"
 if not os.path.exists(clip_l_path):
@@ -105,375 +29,25 @@ esrgan_x4_path = "ComfyUI/models/upscale_models/RealESRGAN_x4.pth"
 if not os.path.exists(esrgan_x4_path):
     hf_hub_download("ai-forever/Real-ESRGAN", "RealESRGAN_x4.pth", local_dir="ComfyUI/models/upscale_models")
 
-# Add ComfyUI to path and import custom nodes
+# ...
+
 add_comfyui_directory_to_sys_path()
 add_extra_model_paths()
-import_custom_nodes()
-
-from nodes import NODE_CLASS_MAPPINGS
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 800px;
-}
-.main-header {
-    text-align: center;
-    margin-bottom: 2rem;
-}
-"""
-
-MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 8192 * 8192
 
-def make_divisible_by_16(size):
-    return ((size // 16) * 16) if (size % 16) < 8 else ((size // 16 + 1) * 16)
+from folder_paths import add_model_folder_path
+add_model_folder_path("checkpoints", "ComfyUI/models/diffusion_models")
 
-def process_input(input_image, upscale_factor):
-    w, h = input_image.size
-    w_original, h_original = w, h
+# ...
 
-    was_resized = False
-
-    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-        gr.Info(f"Requested output image is too large. Resizing input to fit within pixel budget.")
-        target_input_pixels = MAX_PIXEL_BUDGET / (upscale_factor ** 2)
-        scale = (target_input_pixels / (w * h)) ** 0.5
-        new_w = max(16, int(w * scale) // 16 * 16)
-        new_h = max(16, int(h * scale) // 16 * 16)
-        input_image = input_image.resize((new_w, new_h), resample=Image.LANCZOS)
-        was_resized = True
-
-    return input_image, w_original, h_original, was_resized
-
-import requests
-def load_image_from_url(url):
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        return Image.open(response.raw)
-    except Exception as e:
-        raise gr.Error(f"Failed to load image from URL: {e}")
-
-def tensor_to_pil(tensor):
-    tensor = tensor.cpu().clamp(0, 1) * 255
-    img = tensor.numpy().astype(np.uint8)[0]
-    return Image.fromarray(img)
-
-@spaces.GPU(duration=120)
-def enhance_image(
-    image_input,
-    image_url,
-    seed,
-    randomize_seed,
-    num_inference_steps,
-    upscale_factor,
-    denoising_strength,
-    custom_prompt,
-    tile_size,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if image_input is not None:
-        true_input_image = image_input
-    elif image_url:
-        true_input_image = load_image_from_url(image_url)
-    else:
-        raise gr.Error("Please provide an image (upload or URL)")
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    input_image, w_original, h_original, was_resized = process_input(true_input_image, upscale_factor)
-
-    if upscale_factor == 2:
-        upscale_model_name = "RealESRGAN_x2.pth"
-    else:
-        upscale_model_name = "RealESRGAN_x4.pth"
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        input_image.save(tmp.name)
-        image_path = tmp.name
-
-    with torch.inference_mode():
-        dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-        dualcliploader_res = dualcliploader.load_clip(
-            clip_name1="clip_l.safetensors",
-            clip_name2="t5xxl_fp8_e4m3fn.safetensors",
-            type="flux",
+        checkpointloader = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
+        checkpointloader_res = checkpointloader.load_checkpoint(
+            ckpt_name="flux1-dev-fp8.safetensors"
         )
-        clip = get_value_at_index(dualcliploader_res, 0)
-
-        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
-        positive_res = cliptextencode.encode(
-            text=custom_prompt,
-            clip=clip
-        )
-        negative_res = cliptextencode.encode(
-            text="",
-            clip=clip
-        )
-
-        upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
-        upscalemodelloader_res = upscalemodelloader.load_model(
-            model_name=upscale_model_name
-        )
-
-        vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-        vaeloader_res = vaeloader.load_vae(vae_name="ae.safetensors")
-
-        unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-        unetloader_res = unetloader.load_unet(
-            unet_name="flux1-dev-fp8.safetensors", weight_dtype="fp8_e4m3fn"
-        )
-
-        loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
-        loadimage_res = loadimage.load_image(image=os.path.basename(image_path))
-
-        fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
-        fluxguidance_res = fluxguidance.append(
-            guidance=30, conditioning=get_value_at_index(positive_res, 0)
-        )
-
-        ultimatesdupscale = NODE_CLASS_MAPPINGS["UltimateSDUpscale"]()
-        usd_res = ultimatesdupscale.upscale(
-            upscale_by=upscale_factor,
-            seed=seed,
-            steps=num_inference_steps,
-            cfg=1,
-            sampler_name="euler",
-            scheduler="normal",
-            denoise=denoising_strength,
-            mode_type="Linear",
-            tile_width=tile_size,
-            tile_height=tile_size,
-            mask_blur=8,
-            tile_padding=32,
-            seam_fix_mode="None",
-            seam_fix_denoise=1,
-            seam_fix_width=64,
-            seam_fix_mask_blur=8,
-            seam_fix_padding=16,
-            force_uniform_tiles=True,
-            tiled_decode=False,
-            image=get_value_at_index(loadimage_res, 0),
-            model=get_value_at_index(unetloader_res, 0),
-            positive=get_value_at_index(fluxguidance_res, 0),
-            negative=get_value_at_index(negative_res, 0),
-            vae=get_value_at_index(vaeloader_res, 0),
-            upscale_model=get_value_at_index(upscalemodelloader_res, 0),
-        )
-
-        output_tensor = get_value_at_index(usd_res, 0)
-        image = tensor_to_pil(output_tensor)
-
-    os.unlink(image_path)
-
-    target_w, target_h = w_original * upscale_factor, h_original * upscale_factor
-    if image.size != (target_w, target_h):
-        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
-
-    if was_resized:
-        gr.Info(f"Resizing output to target size: {target_w}x{target_h}")
-        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
-
-    resized_input = true_input_image.resize(image.size, resample=Image.LANCZOS)
-
-    return [resized_input, image]
-
-with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX ComfyUI") as demo:
-    gr.HTML("""
-    <div class="main-header">
-        <h1>🎨 AI Image Upscaler (ComfyUI Workflow)</h1>
-        <p>Upload an image or provide a URL to upscale it using FLUX FP8 with ComfyUI Ultimate SD Upscale</p>
-        <p>Using FLUX.1-dev FP8 model</p>
-    </div>
-    """)
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.HTML("<h3>📤 Input</h3>")
-            
-            with gr.Tabs():
-                with gr.TabItem("📁 Upload Image"):
-                    input_image = gr.Image(
-                        label="Upload Image",
-                        type="pil",
-                        height=200
-                    )
-                
-                with gr.TabItem("🔗 Image URL"):
-                    image_url = gr.Textbox(
-                        label="Image URL",
-                        placeholder="https://example.com/image.jpg",
-                        value="https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg"
-                    )
-            
-            gr.HTML("<h3>🎛️ Prompt Settings</h3>")
-            
-            custom_prompt = gr.Textbox(
-                label="Custom Prompt (optional)",
-                placeholder="Enter custom prompt or leave empty",
-                lines=2
-            )
-            
-            gr.HTML("<h3>⚙️ Upscaling Settings</h3>")
-            
-            upscale_factor = gr.Slider(
-                label="Upscale Factor",
-                minimum=1,
-                maximum=4,
-                step=1,
-                value=2,
-                info="How much to upscale the image"
-            )
-            
-            num_inference_steps = gr.Slider(
-                label="Number of Inference Steps",
-                minimum=1,
-                maximum=50,
-                step=1,
-                value=25,
-                info="More steps = better quality but slower"
-            )
-            
-            denoising_strength = gr.Slider(
-                label="Denoising Strength",
-                minimum=0.0,
-                maximum=1.0,
-                step=0.05,
-                value=0.3,
-                info="Controls how much the image is transformed"
-            )
-
-            tile_size = gr.Slider(
-                label="Tile Size",
-                minimum=256,
-                maximum=2048,
-                step=64,
-                value=1024,
-                info="Size of tiles for processing (larger = faster but more memory)"
-            )
-            
-            with gr.Row():
-                randomize_seed = gr.Checkbox(
-                    label="Randomize seed",
-                    value=True
-                )
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=42,
-                    interactive=True
-                )
-            
-            enhance_btn = gr.Button(
-                "🚀 Upscale Image",
-                variant="primary",
-                size="lg"
-            )
-
-        with gr.Column(scale=2):
-            gr.HTML("<h3>📊 Results</h3>")
-            
-            result_slider = ImageSlider(
-                type="pil",
-                interactive=False,
-                height=600,
-                elem_id="result_slider",
-                label=None
-            )
 
-    enhance_btn.click(
-        fn=enhance_image,
-        inputs=[
-            input_image,
-            image_url,
-            seed,
-            randomize_seed,
-            num_inference_steps,
-            upscale_factor,
-            denoising_strength,
-            custom_prompt,
-            tile_size
-        ],
-        outputs=[result_slider]
-    )
-    
-    gr.HTML("""
-    <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
-        <p><strong>Note:</strong> This upscaler uses the Flux.1-dev model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
-    </div>
-    """)
-    
-    gr.HTML("""
-    <style>
-        #result_slider .slider {
-            width: 100% !important;
-            max-width: inherit !important;
-        }
-        #result_slider img {
-            object-fit: contain !important;
-            width: 100% !important;
-            height: auto !important;
-        }
-        #result_slider .gr-button-tool {
-            display: none !important;
-        }
-        #result_slider .gr-button-undo {
-            display: none !important;
-        }
-        #result_slider .gr-button-clear {
-            display: none !important;
-        }
-        #result_slider .badge-container .badge {
-            display: none !important;
-        }
-        #result_slider .badge-container::before {
-            content: "Before";
-            position: absolute;
-            top: 10px;
-            left: 10px;
-            background: rgba(0,0,0,0.5);
-            color: white;
-            padding: 5px;
-            border-radius: 5px;
-            z-index: 10;
-        }
-        #result_slider .badge-container::after {
-            content: "After";
-            position: absolute;
-            top: 10px;
-            right: 10px;
-            background: rgba(0,0,0,0.5);
-            color: white;
-            padding: 5px;
-            border-radius: 5px;
-            z-index: 10;
-        }
-        #result_slider .fullscreen img {
-            object-fit: contain !important;
-            width: 100vw !important;
-            height: 100vh !important;
-            position: absolute;
-            top: 0;
-            left: 0;
-        }
-    </style>
-    """)
-    
-    gr.HTML("""
-    <script>
-        document.addEventListener('DOMContentLoaded', function() {
-            const sliderInput = document.querySelector('#result_slider input[type="range"]');
-            if (sliderInput) {
-                sliderInput.value = 50;
-                sliderInput.dispatchEvent(new Event('input'));
-            }
-        });
-    </script>
-    """)
+        # ...
 
-if __name__ == "__main__":
-    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)
+        ultimatesdupscale_50 = ultimatesdupscale.upscale(
+            # ...
+            model=get_value_at_index(checkpointloader_res, 0),
+            # ...
+        )