Update app.py

app.py

@@ -1,21 +1,33 @@
+import json
+import requests
+from io import BytesIO
 import gradio as gr
 import os
-hf_token = os.environ.get("HF_TOKEN")
+# hf_token = os.environ.get("HF_TOKEN")
 import spaces
-import torch
-from pipeline_bria import BriaPipeline, BriaTransformer2DModel
+# import torch
+# from pipeline_bria import BriaPipeline
 import time
+from PIL import Image
+
+def download_image(url):
+    response = requests.get(url)
+    return Image.open(BytesIO(response.content)).convert("RGB")
+
+hf_token = os.environ.get("HF_TOKEN_API_DEMO") # we get it from a secret env variable, such that it's private
+auth_headers = {"api_token": hf_token}
 
 resolutions = ["1024 1024","1280 768","1344 768","768 1344","768 1280"] 
 
 # Ng
-default_negative_prompt= "Logo,Watermark,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
+default_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,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
 
-transformer = BriaTransformer2DModel.from_pretrained("briaai/BRIA-3.2",subfolder='transformer',torch_dtype=torch.bfloat16)
-pipe = BriaPipeline.from_pretrained("briaai/BRIA-3.1", transformer=transformer, torch_dtype=torch.bfloat16,trust_remote_code=True)
-pipe.to(device="cuda")
+# Load pipeline
+# trust_remote_code = True - allows loading a transformer which is not present at the transformers library(from transformer/bria_transformer.py)
+# pipe = BriaPipeline.from_pretrained("briaai/BRIA-3.0-TOUCAN", torch_dtype=torch.bfloat16,trust_remote_code=True)
+# pipe.to(device="cuda")
 
-@spaces.GPU(enable_queue=True)
+# @spaces.GPU(enable_queue=True)
 def infer(prompt,negative_prompt,seed,resolution):
     print(f"""
     —/n
@@ -30,13 +42,34 @@ def infer(prompt,negative_prompt,seed,resolution):
     else:
         try:
             seed=int(seed)
-            generator = torch.Generator("cuda").manual_seed(seed)
+            # generator = torch.Generator("cuda").manual_seed(seed)
         except:
             generator=None
 
     w,h = resolution.split()
     w,h = int(w),int(h)
-    image = pipe(prompt,num_inference_steps=30, negative_prompt=negative_prompt,generator=generator,width=w,height=h).images[0]
+    # image = pipe(prompt,num_inference_steps=30, negative_prompt=negative_prompt,generator=generator,width=w,height=h).images[0]
+    url = "http://engine.prod.bria-api.com/v1/text-to-image/base/3.2"
+    
+    payload = json.dumps({
+    "prompt": prompt,
+    "num_results": 1,
+    "sync": True,
+    "prompt_enhancement": False,
+    "debias": False,
+    "fast": False,
+    "model_influence": 0.0000001,
+    "include_generation_prefix": False,
+    "negative_prompt": negative_prompt,
+    "num_inference_steps": 30,
+    "seed": seed
+    })
+    response = requests.request("POST", url, headers=auth_headers, data=payload)
+    print('1',response)
+    response = response.json()
+    print('2',response)
+    res_image = download_image(response["result"][0]['urls'][0])
+
     print(f'gen time is {time.time()-t} secs')
     
     # Future
@@ -44,7 +77,7 @@ def infer(prompt,negative_prompt,seed,resolution):
     # if nsfw:
     #     raise gr.Error("Generated image is NSFW")
     
-    return image
+    return res_image
 
 css = """
 #col-container{
@@ -54,7 +87,7 @@ css = """
 """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("## BRIA 3.2")
+        gr.Markdown("## BRIA-3.2")
         gr.HTML('''
           <p style="margin-bottom: 10px; font-size: 94%">
             This is a demo for 
@@ -70,7 +103,7 @@ with gr.Blocks(css=css) as demo:
         ''')
         with gr.Group():
             with gr.Column():
-                prompt_in = gr.Textbox(label="Prompt", value="""photo of mystical dragon eating sushi, text bubble says "Sushi Time".""")
+                prompt_in = gr.Textbox(label="Prompt", value="A smiling man with wavy brown hair and a trimmed beard")
                 resolution = gr.Dropdown(value=resolutions[0], show_label=True, label="Resolution", choices=resolutions)
                 seed = gr.Textbox(label="Seed", value=-1)
                 negative_prompt = gr.Textbox(label="Negative Prompt", value=default_negative_prompt)