Update app.py

app.py

@@ -25,6 +25,17 @@ login(token=hf_token)
 
 MAX_SEED = np.iinfo(np.int32).max
 
+try:
+    local_dir = os.path.dirname(__file__)
+except:
+    local_dir = '.'
+    
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='pipeline_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='transformer_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='bria_utils.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt-ControlNet-Union", filename='pipeline_bria_controlnet.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt-ControlNet-Union", filename='controlnet_bria.py', local_dir=local_dir)
+
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -67,10 +78,9 @@ import PIL.Image as Image
 
 base_model = 'briaai/BRIA-4B-Adapt'
 controlnet_model = 'briaai/BRIA-4B-Adapt-ControlNet-Union'
-
 controlnet = BriaControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
 pipe = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
-pipe.to("cuda")
+pipe = pipeline.to(device="cuda", dtype=torch.bfloat16)
 
 mode_mapping = {
     "depth": 0,
@@ -172,6 +182,7 @@ def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mo
         guidance_scale=guidance_scale,
         generator=torch.manual_seed(seed),
         max_sequence_length=128,
+        negative_prompt="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate"
     ).images[0]
 
     torch.cuda.empty_cache()