Update app.py

app.py

@@ -5,8 +5,7 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image, FluxControlNetModel
-from diffusers.pipelines import FluxControlNetPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from diffusers.utils import load_image
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
@@ -21,7 +20,7 @@ import numpy as np
 import warnings
 
 
-huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+huggingface_token = os.getenv("HF_TOKEN")
 
 
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
@@ -61,23 +60,6 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     torch_dtype=dtype
 ).to(device)
 
-# Upscale을 위한 ControlNet 설정
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-).to(device)
-
-# Upscale 파이프라인 설정 (기존 pipe 재사용)
-pipe_upscale = FluxControlNetPipeline(
-    vae=pipe.vae,
-    text_encoder=pipe.text_encoder,
-    text_encoder_2=pipe.text_encoder_2,
-    tokenizer=pipe.tokenizer,
-    tokenizer_2=pipe.tokenizer_2,
-    transformer=pipe.transformer,
-    scheduler=pipe.scheduler,
-    controlnet=controlnet
-).to(device)
-
 MAX_SEED = 2**32 - 1
 MAX_PIXEL_BUDGET = 1024 * 1024
 
@@ -556,109 +538,6 @@ css = '''
 footer {visibility: hidden;}
 '''
 
-# 업스케일 관련 함수 추가
-def process_input(input_image, upscale_factor, **kwargs):
-    w, h = input_image.size
-    w_original, h_original = w, h
-    aspect_ratio = w / h
-
-    was_resized = False
-
-    max_size = int(np.sqrt(MAX_PIXEL_BUDGET / (upscale_factor ** 2)))
-    if w > max_size or h > max_size:
-        if w > h:
-            w_new = max_size
-            h_new = int(w_new / aspect_ratio)
-        else:
-            h_new = max_size
-            w_new = int(h_new * aspect_ratio)
-        
-        input_image = input_image.resize((w_new, h_new), Image.LANCZOS)
-        was_resized = True
-        gr.Info(f"Input image resized to {w_new}x{h_new} to fit within pixel budget after upscaling.")
-
-    # resize to multiple of 8
-    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
-    
-from PIL import Image
-import numpy as np
-
-@spaces.GPU
-def infer_upscale(
-    seed,
-    randomize_seed,
-    input_image,
-    num_inference_steps,
-    upscale_factor,
-    controlnet_conditioning_scale,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if input_image is None:
-        return None, seed, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(visible=True, value="Please upload an image for upscaling.")
-
-    try:
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-        
-        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), Image.LANCZOS)
-
-        generator = torch.Generator(device=device).manual_seed(seed)
-
-        gr.Info("Upscaling image...")
-        # 모든 텐서를 동일한 디바이스로 이동
-        pipe_upscale.to(device)
-        
-        # Ensure the image is in RGB format
-        if control_image.mode != 'RGB':
-            control_image = control_image.convert('RGB')
-        
-        # Convert to tensor and add batch dimension
-        control_image = torch.from_numpy(np.array(control_image)).permute(2, 0, 1).float().unsqueeze(0).to(device) / 255.0
-        
-        with torch.no_grad():
-            image = pipe_upscale(
-                prompt="",
-                control_image=control_image,
-                controlnet_conditioning_scale=controlnet_conditioning_scale,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=3.5,
-                generator=generator,
-            ).images[0]
-
-        # Convert the image back to PIL Image
-        if isinstance(image, torch.Tensor):
-            image = image.cpu().permute(1, 2, 0).numpy()
-        
-        # Ensure the image data is in the correct range
-        image = np.clip(image * 255, 0, 255).astype(np.uint8)
-        image = Image.fromarray(image)
-
-        if was_resized:
-            gr.Info(
-                f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
-            )
-            image = image.resize((w_original * upscale_factor, h_original * upscale_factor), Image.LANCZOS)
-
-        return image, seed, num_inference_steps, upscale_factor, controlnet_conditioning_scale, gr.update(), gr.update(visible=False)
-    except Exception as e:
-        print(f"Error in infer_upscale: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return None, seed, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(visible=True, value=f"Error: {str(e)}")
-        
-def check_upscale_input(input_image, *args):
-    if input_image is None:
-        return gr.update(interactive=False), *args, gr.update(visible=True, value="Please upload an image for upscaling.")
-    return gr.update(interactive=True), *args, gr.update(visible=False)
-    
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as app:
     loras_state = gr.State(loras)
     selected_indices = gr.State([])
@@ -742,49 +621,6 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
 
-    # 업스케일 관련 UI 추가
-    with gr.Row():
-        upscale_button = gr.Button("Upscale", interactive=False)
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            upscale_input = gr.Image(label="Input Image for Upscaling", type="pil")
-        with gr.Column(scale=1):
-            upscale_steps = gr.Slider(
-                label="Number of Inference Steps for Upscaling",
-                minimum=8,
-                maximum=50,
-                step=1,
-                value=28,
-            )
-            upscale_factor = gr.Slider(
-                label="Upscale Factor",
-                minimum=1,
-                maximum=4,
-                step=1,
-                value=4,
-            )
-            controlnet_conditioning_scale = gr.Slider(
-                label="Controlnet Conditioning Scale",
-                minimum=0.1,
-                maximum=1.0,
-                step=0.05,
-                value=0.5,  # 기본값을 0.5로 낮춤
-            )            
-            upscale_seed = gr.Slider(
-                label="Seed for Upscaling",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=42,
-            )
-            upscale_randomize_seed = gr.Checkbox(label="Randomize seed for Upscaling", value=True)
-            upscale_error = gr.Markdown(visible=False, value="Please provide an input image for upscaling.")
-    
-    with gr.Row():
-        upscale_result = gr.Image(label="Upscaled Image", type="pil")
-        upscale_seed_output = gr.Number(label="Seed Used", precision=0)
-
     gallery.select(
         update_selection,
         inputs=[selected_indices, loras_state, width, height],
@@ -809,8 +645,6 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         outputs=[selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3]
     )
 
- 
-
     randomize_button.click(
         randomize_loras,
         inputs=[selected_indices, loras_state],
@@ -840,46 +674,6 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         outputs=history_gallery,
     )
 
-    upscale_input.upload(
-        lambda x: gr.update(interactive=x is not None),
-        inputs=[upscale_input],
-        outputs=[upscale_button]
-    )
-    
-    upscale_error = gr.Markdown(visible=False, value="")
-
-    upscale_button.click(
-        infer_upscale,
-        inputs=[
-            upscale_seed,
-            upscale_randomize_seed,
-            upscale_input,
-            upscale_steps,
-            upscale_factor,
-            controlnet_conditioning_scale,
-        ],
-        outputs=[
-            upscale_result,
-            upscale_seed_output,
-            upscale_steps,
-            upscale_factor,
-            controlnet_conditioning_scale,
-            upscale_randomize_seed,
-            upscale_error
-        ],
-    ).then(
-        infer_upscale,
-        inputs=[
-        upscale_seed,
-            upscale_randomize_seed,
-            upscale_input,
-            upscale_steps,
-            upscale_factor,
-            controlnet_conditioning_scale,
-        ],
-        outputs=[upscale_result, upscale_seed_output]
-    )
-
 if __name__ == "__main__":
     app.queue(max_size=20)
     app.launch(debug=True)