Update app.py

app.py

@@ -1,17 +1,12 @@
 import gradio as gr
 import numpy as np
-#import random
-import spaces #[uncomment to use ZeroGPU]
-#from diffusers import DiffusionPipeline
 import torch
 
-#SAM
-from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
+import spaces 
+
 #Flux
 from diffusers import FluxFillPipeline
 
-#import cv2
-#import matplotlib.pyplot as plt
 from PIL import Image
 import os
 import gc
@@ -20,74 +15,18 @@ import glob
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# GDINO_MODEL_NAME="IDEA-Research/grounding-dino-tiny" 
-# SAM_MODEL_NAME="facebook/sam-vit-base" 
-
-# GDINO=pipeline(model=GDINO_MODEL_NAME, task="zero-shot-object-detection", device=DEVICE)
-# SAM=AutoModelForMaskGeneration.from_pretrained(SAM_MODEL_NAME).to(DEVICE)
-# SAM_PROCESSOR=AutoProcessor.from_pretrained(SAM_MODEL_NAME)
-
-
 SD_PIPLINE = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-
-#except Exception as ex:
-    #print(ex)
-
-
-# SD_MODEL="black-forest-labs/FLUX.1-dev" #"diffusers/stable-diffusion-xl-1.0-inpainting-0.1" 
-# SD_PIPLINE = AutoPipelineForInpainting.from_pretrained(SD_MODEL, torch_dtype=torch.float16).to(DEVICE)
-# IP_ADAPTER='XLabs-AI/flux-ip-adapter-v2' #"h94/IP-Adapter" 
-# #SUB_FOLDER="sdxl_models"  
-# IP_WEIGHT_NAME='ip_adapter.safetensors' #"ip-adapter_sdxl.bin" 
-# SD_PIPLINE.load_ip_adapter(IP_ADAPTER, weight_name=IP_WEIGHT_NAME) #subfolder=SUB_FOLDER, weight_name=IP_WEIGHT_NAME)
-# IP_SCALE=0.6 
-# SD_PIPLINE.set_ip_adapter_scale(IP_SCALE)
-
-GEN_STEPS=40 #20
-
-
-def refine_masks(masks: torch.BoolTensor)->np.array:    
-    masks = masks.permute(0, 2, 3, 1)
-    masks = masks.float()#.mean(axis=-1)
-    return masks.cpu().numpy()
-
-
-def get_boxes(detections:list)-> list:
-    boxes = []
-    for det in detections:
-        boxes.append([det['box']['xmin'], det['box']['ymin'], 
-                      det['box']['xmax'], det['box']['ymax']])
-    return [boxes]
-
-
-def get_mask(img:Image, prompt:str, d_model:pipeline, s_model:AutoModelForMaskGeneration,
-             s_processor:AutoProcessor, device:str, threshold:float = 0.3)-> np.array:
-    
-    labels = [label if label.endswith(".") else label+"." for label in ['face', prompt]] 
-    dets=d_model(img, candidate_labels=labels, threshold=threshold)    
-
-    boxes = get_boxes(dets)
-    inputs=s_processor(images=img, input_boxes=boxes, return_tensors="pt").to(DEVICE)
-    outputs = s_model(**inputs)
-    
-    masks = s_processor.post_process_masks(
-        masks=outputs.pred_masks,
-        original_sizes=inputs.original_sizes,
-        reshaped_input_sizes=inputs.reshaped_input_sizes
-    )[0] 
-
-    return refine_masks(masks) 
+GEN_STEPS=35 
 
 
-def generate_result(model_img:str, #cloth_img:str,
-                    cloth_mask,#: np.array, 
-                    prompt:str, sd_pipline,#:FluxFillPipeline, 
+def generate_result(model_img:Image, 
+                    cloth_mask:Image,
+                    prompt:str, sd_pipline: FluxFillPipeline,
                     n_steps:int=100)->Image:
     
-    width, height = model_img.size
+    width, height = model_img.size   
     
-    #cloth_mask=Image.fromarray(masks[1].astype(np.uint8)) #np.array(masks[1],dtype=np.float32)
     generator = torch.Generator(device="cpu").manual_seed(4)  
     
     images = sd_pipline(
@@ -104,9 +43,8 @@ def generate_result(model_img:str, #cloth_img:str,
 
 
 @spaces.GPU
-def run(model_img:Image,mask:Image, cloth_class:str, close_description:str)->Image: #cloth_img:Image, cloth_class:str, close_description:str)->Image:
-    #masks = get_mask(model_img, cloth_class, GDINO, SAM, SAM_PROCESSOR, DEVICE) #GSAM2)
-    result = generate_result(model_img,mask, close_description, SD_PIPLINE, GEN_STEPS) #cloth_img, masks, close_description, SD_PIPLINE, GEN_STEPS)
+def run(model_img:Image, mask:Image, prompt:str)->Image:    
+    result = generate_result(model_img, mask, prompt, SD_PIPLINE, GEN_STEPS) 
     gc.collect()
     torch.cuda.empty_cache()
     return result
@@ -114,12 +52,12 @@ def run(model_img:Image,mask:Image, cloth_class:str, close_description:str)->Ima
 
 gr.Interface(
     run,
-    title = 'Virtual Try-On',
+    title = 'Try-On Flux',
     inputs=[
         gr.Image(sources = 'upload', label='Model image', type = 'pil'),
         gr.Image(sources = 'upload', label='Mask', type = 'pil'),
-        gr.Textbox(label = 'Cloth class'),
-        gr.Textbox(label = 'Close description')
+        #gr.Textbox(label = 'Cloth class'),
+        gr.Textbox(label = 'Prompt')
     ],
     outputs = [
         gr.Image()