Update app.py

app.py

@@ -12,6 +12,11 @@ from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
 
+# 메모리 관리를 위한 gc 추가
+import gc
+gc.collect()
+torch.cuda.empty_cache()
+
 css = """
 #col-container {
     margin: 0 auto;
@@ -19,97 +24,96 @@ css = """
 }
 """
 
-# Device and dtype setup with lower precision
+# Device setup with minimal memory usage
 if torch.cuda.is_available():
     power_device = "GPU"
     device = "cuda" 
-    dtype = torch.float16  # Changed to float16 for less memory usage
+    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
 
-# Reduce CUDA memory usage
-torch.cuda.empty_cache()
-if torch.cuda.is_available():
-    torch.cuda.set_per_process_memory_fraction(0.7)  # Use only 70% of GPU memory
-
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
+# Minimal model configuration
+model_config = {
+    "low_cpu_mem_usage": True,
+    "torch_dtype": dtype,
+    "use_safetensors": True,
+    "variant": "fp16",  # Use fp16 variant if available
+}
+
 model_path = snapshot_download(
     repo_id="black-forest-labs/FLUX.1-dev",
     repo_type="model",
-    ignore_patterns=["*.md", "*..gitattributes"],
+    ignore_patterns=["*.md", "*..gitattributes", "*.bin"],  # Ignore unnecessary files
     local_dir="FLUX.1-dev", 
     token=huggingface_token,
 )
 
-# Load pipeline with more memory optimizations
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler",
-    torch_dtype=dtype,
-    low_cpu_mem_usage=True,
-    use_safetensors=True
-).to(device)
-
-pipe = FluxControlNetPipeline.from_pretrained(
-    model_path,
-    controlnet=controlnet,
-    torch_dtype=dtype,
-    low_cpu_mem_usage=True,
-    use_safetensors=True
-)
+# 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
+    )
 
-# Enable all possible memory optimizations
-pipe.enable_model_cpu_offload()
-pipe.enable_attention_slicing(1)
-pipe.enable_sequential_cpu_offload()
-pipe.enable_vae_slicing()
+    # 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
 
-# Further reduce memory usage
+# Extremely reduced parameters
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 256 * 256  # Further reduced from 512 * 512
+MAX_PIXEL_BUDGET = 128 * 128  # Extremely reduced from 256 * 256
 
 def check_resources():
     if torch.cuda.is_available():
-        gpu_memory = torch.cuda.get_device_properties(0).total_memory
         memory_allocated = torch.cuda.memory_allocated(0)
-        if memory_allocated/gpu_memory > 0.8:  # 80% threshold
-            return False
+        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):
-    # Convert image to RGB mode to ensure compatibility
     input_image = input_image.convert('RGB')
     
+    # Reduce image size more aggressively
     w, h = input_image.size
-    w_original, h_original = w, h
-    aspect_ratio = w / h
-
-    was_resized = False
-
-    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-        warnings.warn(
-            f"Requested output image is too large. Resizing..."
-        )
-        gr.Info(
-            f"Resizing input image to fit memory constraints..."
-        )
-        input_image = input_image.resize(
-            (
-                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
-                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
-            ),
-            Image.LANCZOS
-        )
-        was_resized = True
-
-    # resize to multiple of 8
+    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
-
-    return input_image.resize((w, h)), w_original, h_original, was_resized
+    
+    return input_image.resize((w, h)), w, h, True
 
 @spaces.GPU
 def infer(
@@ -122,55 +126,32 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     try:
-        if not check_resources():
-            gr.Warning("System resources are running low. Try reducing parameters.")
-            return None
-
-        # Clear CUDA cache before processing
-        if device == "cuda":
-            torch.cuda.empty_cache()
-
+        gc.collect()
+        torch.cuda.empty_cache()
+        
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
             
-        true_input_image = input_image
-        input_image, w_original, h_original, was_resized = process_input(
-            input_image, upscale_factor
-        )
-
-        # rescale with upscale factor
-        w, h = input_image.size
-        control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-        generator = torch.Generator().manual_seed(seed)
-
-        gr.Info("Upscaling image...")
-        with torch.inference_mode():  # Use inference mode to save memory
+        input_image, w, h, _ = process_input(input_image, upscale_factor)
+        
+        with torch.inference_mode():
+            generator = torch.Generator().manual_seed(seed)
             image = pipe(
                 prompt="",
-                control_image=control_image,
+                control_image=input_image,
                 controlnet_conditioning_scale=controlnet_conditioning_scale,
                 num_inference_steps=num_inference_steps,
-                guidance_scale=3.5,
-                height=control_image.size[1],
-                width=control_image.size[0],
+                guidance_scale=2.0,  # Reduced from 3.5
+                height=h,
+                width=w,
                 generator=generator,
             ).images[0]
+            
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        return [input_image, image, seed]
 
-        if was_resized:
-            gr.Info(
-                f"Resizing output image to final size..."
-            )
-
-        # resize to target desired size
-        image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
-        return [true_input_image, image, seed]
-
-    except RuntimeError as e:
-        if "out of memory" in str(e):
-            gr.Warning("Not enough GPU memory. Try reducing the upscale factor or image size.")
-            return None
-        raise e
     except Exception as e:
         gr.Error(f"An error occurred: {str(e)}")
         return None
@@ -184,25 +165,25 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
             input_im = gr.Image(label="Input Image", type="pil")
         with gr.Column(scale=1):
             num_inference_steps = gr.Slider(
-                label="Number of Inference Steps",
-                minimum=8,
-                maximum=30,  # Reduced from 50
+                label="Steps",
+                minimum=1,
+                maximum=20,  # Reduced from 30
                 step=1,
-                value=20,    # Reduced from 28
+                value=10,    # Reduced from 20
             )
             upscale_factor = gr.Slider(
-                label="Upscale Factor",
+                label="Scale",
                 minimum=1,
-                maximum=2,
+                maximum=1,   # Fixed at 1
                 step=1,
-                value=1,    # Reduced default
+                value=1,
             )
             controlnet_conditioning_scale = gr.Slider(
-                label="Controlnet Conditioning Scale",
+                label="Control Scale",
                 minimum=0.1,
-                maximum=1.0,  # Reduced from 1.5
+                maximum=0.5,  # Reduced from 1.0
                 step=0.1,
-                value=0.5,   # Reduced from 0.6
+                value=0.3,   # Reduced from 0.5
             )
             seed = gr.Slider(
                 label="Seed",
@@ -211,18 +192,17 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                 step=1,
                 value=42,
             )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            randomize_seed = gr.Checkbox(label="Random Seed", value=True)
 
     with gr.Row():
-        result = ImageSlider(label="Input / Output", type="pil", interactive=True)
+        result = ImageSlider(label="Result", type="pil", interactive=True)
 
     current_dir = os.path.dirname(os.path.abspath(__file__))
     
     examples = gr.Examples(
         examples=[
-            [42, False, os.path.join(current_dir, "z1.webp"), 20, 1, 0.5],  # Reduced parameters
-            [42, False, os.path.join(current_dir, "z2.webp"), 20, 1, 0.5],  # Reduced parameters
+            [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],
         ],
         inputs=[
             seed,
@@ -234,7 +214,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         ],
         fn=infer,
         outputs=result,
-        cache_examples="lazy",
+        cache_examples=False,  # Disable caching
     )
 
     gr.on(
@@ -252,11 +232,13 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         show_api=False,
     )
 
-# Launch with minimal memory usage
+# Launch with minimal resources
 demo.queue(max_size=1).launch(
     share=False,
     debug=True,
     show_error=True,
     max_threads=1,
-    enable_queue=True
+    enable_queue=True,
+    cache_examples=False,
+    quiet=True,
 )