Update app.py

app.py

@@ -4,18 +4,22 @@ import warnings
 import os
 import gradio as gr
 import numpy as np
-import spaces
 import torch
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
 from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
-
-# 메모리 관리를 위한 gc 추가
 import gc
+
+# Force CPU usage
+device = "cpu"
+dtype = torch.float32
+
+# Clear memory
 gc.collect()
-torch.cuda.empty_cache()
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
 
 css = """
 #col-container {
@@ -24,98 +28,61 @@ css = """
 }
 """
 
-# Device setup with minimal memory usage
-if torch.cuda.is_available():
-    power_device = "GPU"
-    device = "cuda" 
-    dtype = torch.float16  # Use float16 for minimum memory
-    # Set CUDA memory fraction to 50%
-    torch.cuda.set_per_process_memory_fraction(0.5)
-else:
-    power_device = "CPU"
-    device = "cpu"
-    dtype = torch.float32
-
-huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+huggingface_token = os.getenv("HF_TOKEN")
 
 # Minimal model configuration
 model_config = {
     "low_cpu_mem_usage": True,
     "torch_dtype": dtype,
     "use_safetensors": True,
-    "variant": "fp16",  # Use fp16 variant if available
+    "device_map": "cpu"
 }
 
 model_path = snapshot_download(
     repo_id="black-forest-labs/FLUX.1-dev",
     repo_type="model",
-    ignore_patterns=["*.md", "*..gitattributes", "*.bin"],  # Ignore unnecessary files
+    ignore_patterns=["*.md", "*..gitattributes", "*.bin"],
     local_dir="FLUX.1-dev", 
     token=huggingface_token,
 )
 
-# Load models with minimal configuration
-try:
-    controlnet = FluxControlNetModel.from_pretrained(
-        "jasperai/Flux.1-dev-Controlnet-Upscaler",
-        **model_config
-    ).to(device)
-
-    pipe = FluxControlNetPipeline.from_pretrained(
-        model_path,
-        controlnet=controlnet,
-        **model_config
-    )
+# Load models on CPU
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    **model_config
+)
 
-    # Enable all memory optimizations
-    pipe.enable_model_cpu_offload()
-    pipe.enable_attention_slicing(1)
-    pipe.enable_sequential_cpu_offload()
-    pipe.enable_vae_slicing()
-    
-    # Clear memory after loading
-    gc.collect()
-    torch.cuda.empty_cache()
-    
-except Exception as e:
-    print(f"Error loading models: {e}")
-    raise
+pipe = FluxControlNetPipeline.from_pretrained(
+    model_path,
+    controlnet=controlnet,
+    **model_config
+)
+
+# Enable optimizations
+pipe.enable_attention_slicing(1)
+pipe.enable_vae_slicing()
 
-# Extremely reduced parameters
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 128 * 128  # Extremely reduced from 256 * 256
-
-def check_resources():
-    if torch.cuda.is_available():
-        memory_allocated = torch.cuda.memory_allocated(0)
-        memory_reserved = torch.cuda.memory_reserved(0)
-        if memory_allocated/memory_reserved > 0.7:  # 70% threshold
-            gc.collect()
-            torch.cuda.empty_cache()
-    return True
-
-def process_input(input_image, upscale_factor, **kwargs):
+MAX_PIXEL_BUDGET = 64 * 64  # Extremely reduced
+
+def process_input(input_image, upscale_factor):
     input_image = input_image.convert('RGB')
     
-    # Reduce image size more aggressively
+    # Aggressive size reduction
     w, h = input_image.size
     max_size = int(np.sqrt(MAX_PIXEL_BUDGET))
-    if w > max_size or h > max_size:
-        if w > h:
-            new_w = max_size
-            new_h = int(h * max_size / w)
-        else:
-            new_h = max_size
-            new_w = int(w * max_size / h)
-        input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
     
-    w, h = input_image.size
-    w = w - w % 8
-    h = h - h % 8
+    # Resize to very small size
+    new_w = min(w, max_size)
+    new_h = min(h, max_size)
+    input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
+    
+    # Ensure dimensions are multiples of 8
+    w = new_w - new_w % 8
+    h = new_h - new_h % 8
     
-    return input_image.resize((w, h)), w, h, True
+    return input_image.resize((w, h)), w, h
 
-@spaces.GPU
 def infer(
     seed,
     randomize_seed,
@@ -127,12 +94,11 @@ def infer(
 ):
     try:
         gc.collect()
-        torch.cuda.empty_cache()
         
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
             
-        input_image, w, h, _ = process_input(input_image, upscale_factor)
+        input_image, w, h = process_input(input_image, upscale_factor)
         
         with torch.inference_mode():
             generator = torch.Generator().manual_seed(seed)
@@ -141,19 +107,18 @@ def infer(
                 control_image=input_image,
                 controlnet_conditioning_scale=controlnet_conditioning_scale,
                 num_inference_steps=num_inference_steps,
-                guidance_scale=2.0,  # Reduced from 3.5
+                guidance_scale=1.5,
                 height=h,
                 width=w,
                 generator=generator,
             ).images[0]
             
         gc.collect()
-        torch.cuda.empty_cache()
         
         return [input_image, image, seed]
 
     except Exception as e:
-        gr.Error(f"An error occurred: {str(e)}")
+        gr.Error(f"Error: {str(e)}")
         return None
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
@@ -167,23 +132,23 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
             num_inference_steps = gr.Slider(
                 label="Steps",
                 minimum=1,
-                maximum=20,  # Reduced from 30
+                maximum=10,
                 step=1,
-                value=10,    # Reduced from 20
+                value=5,
             )
             upscale_factor = gr.Slider(
                 label="Scale",
                 minimum=1,
-                maximum=1,   # Fixed at 1
+                maximum=1,
                 step=1,
                 value=1,
             )
             controlnet_conditioning_scale = gr.Slider(
                 label="Control Scale",
                 minimum=0.1,
-                maximum=0.5,  # Reduced from 1.0
+                maximum=0.3,
                 step=0.1,
-                value=0.3,   # Reduced from 0.5
+                value=0.2,
             )
             seed = gr.Slider(
                 label="Seed",
@@ -201,8 +166,8 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     
     examples = gr.Examples(
         examples=[
-            [42, False, os.path.join(current_dir, "z1.webp"), 10, 1, 0.3],
-            [42, False, os.path.join(current_dir, "z2.webp"), 10, 1, 0.3],
+            [42, False, os.path.join(current_dir, "z1.webp"), 5, 1, 0.2],
+            [42, False, os.path.join(current_dir, "z2.webp"), 5, 1, 0.2],
         ],
         inputs=[
             seed,
@@ -214,7 +179,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         ],
         fn=infer,
         outputs=result,
-        cache_examples=False,  # Disable caching
+        cache_examples=False,
     )
 
     gr.on(
@@ -232,7 +197,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         show_api=False,
     )
 
-# Launch with minimal resources
+# Minimal launch configuration
 demo.queue(max_size=1).launch(
     share=False,
     debug=True,