Update app.py

app.py

@@ -75,11 +75,55 @@ model = AutoModelForCausalLM.from_pretrained(model_id, token= token,
 model = accelerator.prepare(model)
 
 
-# device_map = infer_auto_device_map(model, max_memory={0: "79GB", "cpu":"65GB" })
 
-# Load the model with the inferred device map
-# model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, no_split_module_classes=["GPTJBlock"])
-# model.half()
+
+###################################################    BG REMOVER ###################################################
+
+
+import gradio as gr
+from gradio_imageslider import ImageSlider
+from loadimg import load_img
+import spaces
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to("cuda")
+
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+    
+
+
+import base64
+from io import BytesIO
+from PIL import Image
+
+def convert_image_to_base64(image):
+    """
+    Convert a PIL Image with alpha channel to a base64-encoded string.
+    """
+    # Save the image into a BytesIO buffer
+    img_byte_array = BytesIO()
+    image.save(img_byte_array, format="PNG")  # Use PNG for transparency
+    img_byte_array.seek(0)  # Reset the pointer to the beginning
+
+    # Encode the image bytes to base64
+    base64_str = base64.b64encode(img_byte_array.getvalue()).decode("utf-8")
+    return base64_str
+
+
 
 import json
 
@@ -106,27 +150,48 @@ def respond(
     
     messages= json_obj
 
-    input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(accelerator.device)
-    input_ids2 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt") #.to('cuda')
-    print(f"Converted input_ids dtype: {input_ids.dtype}")
-    input_str= str(input_ids2)
-    print('input str = ', input_str)
-
-    with torch.no_grad():
-        gen_tokens = model.generate(
-    input_ids, 
-    max_new_tokens=max_tokens, 
-    # do_sample=True, 
-    temperature=temperature,
-    )
-
-    gen_text = tokenizer.decode(gen_tokens[0])
-    print(gen_text)
-    gen_text= gen_text.replace(input_str,'')
-    gen_text= gen_text.replace('<|eot_id|>','')
+    try:
+        image= json_obj['image']
+        image = load_img(image, output_type="pil")
+        image = im.convert("RGB")
+        
+        image_size = image.size
     
-    yield gen_text
-   
+        input_images = transform_image(image).unsqueeze(0).to("cuda")
+        # Prediction
+        with torch.no_grad():
+            preds = birefnet(input_images)[-1].sigmoid().cpu()
+        pred = preds[0].squeeze()
+        pred_pil = transforms.ToPILImage()(pred)
+        mask = pred_pil.resize(image_size)
+        image.putalpha(mask)
+        return convert_image_to_base64(image)
+        
+
+    except Exception as e:
+        print("using llama 8b intrcuxt ",e)
+
+        input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(accelerator.device)
+        input_ids2 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt") #.to('cuda')
+        print(f"Converted input_ids dtype: {input_ids.dtype}")
+        input_str= str(input_ids2)
+        print('input str = ', input_str)
+    
+        with torch.no_grad():
+            gen_tokens = model.generate(
+        input_ids, 
+        max_new_tokens=max_tokens, 
+        # do_sample=True, 
+        temperature=temperature,
+        )
+    
+        gen_text = tokenizer.decode(gen_tokens[0])
+        print(gen_text)
+        gen_text= gen_text.replace(input_str,'')
+        gen_text= gen_text.replace('<|eot_id|>','')
+        
+        yield gen_text
+       
   
 #     messages = [
 #     # {"role": "user", "content": "What is your favourite condiment?"},