Upload 2 files

app.py

@@ -0,0 +1,110 @@
+import gradio as gr
+import os
+hf_token = os.environ.get("HF_TOKEN")
+import spaces
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+import torch
+import time
+
+class Dummy():
+    pass
+
+# Load pipeline
+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"
+model_id = "briaai/BRIA-2.2"
+scheduler = EulerAncestralDiscreteScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                num_train_timesteps=1000,
+                steps_offset=1
+            )
+pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16,scheduler=scheduler).to("cuda")
+
+print("Optimizing BRIA-2.2 - this could take a while")
+t=time.time()
+pipe.unet = torch.compile(
+    pipe.unet, mode="reduce-overhead", fullgraph=True # 600 secs compilation
+)
+with torch.no_grad():
+    outputs = pipe(
+        prompt="an apple",
+        num_inference_steps=30,
+    )
+
+    # This will avoid future compilations on different shapes
+    unet_compiled = torch._dynamo.run(pipe.unet)
+    unet_compiled.config=pipe.unet.config
+    unet_compiled.add_embedding = Dummy()
+    unet_compiled.add_embedding.linear_1 = Dummy()
+    unet_compiled.add_embedding.linear_1.in_features = pipe.unet.add_embedding.linear_1.in_features
+    pipe.unet = unet_compiled
+
+print(f"Optimizing finished successfully after {time.time()-t} secs")
+
+@spaces.GPU(enable_queue=True)
+def infer(prompt):
+    print(f"""
+    —/n
+    {prompt}
+    """)
+    
+    # generator = torch.Generator("cuda").manual_seed(555)
+    t=time.time()
+    image = pipe(prompt,num_inference_steps=30, negative_prompt=default_negative_prompt).images[0]
+    print(f'gen time is {time.time()-t} secs')
+    
+    # Future
+    # Add amound of steps
+    # if nsfw:
+    #     raise gr.Error("Generated image is NSFW")
+    
+    return image
+
+css = """
+#col-container{
+    margin: 0 auto;
+    max-width: 580px;
+}
+"""
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.HTML("""
+        <h2 style="text-align: center;">
+            BRIA-2.2
+        </h2>
+        """)
+        with gr.Group():
+            with gr.Column():
+                prompt_in = gr.Textbox(label="Prompt", value="A red colored sports car")
+                submit_btn = gr.Button("Generate")
+        result = gr.Image(label="BRIA-2.2 Result")
+
+        # gr.Examples(
+        #     examples = [ 
+        #         "Dragon, digital art, by Greg Rutkowski",
+        #         "Armored knight holding sword",
+        #         "A flat roof villa near a river with black walls and huge windows",
+        #         "A calm and peaceful office",
+        #         "Pirate guinea pig"
+        #     ],
+        #     fn = infer, 
+        #     inputs = [
+        #         prompt_in
+        #     ],
+        #     outputs = [
+        #         result
+        #     ]
+        # )
+
+    submit_btn.click(
+        fn = infer,
+        inputs = [
+            prompt_in
+        ],
+        outputs = [
+            result
+        ]
+    )
+
+demo.queue().launch(show_api=False)

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+diffusers
+torch
+torchvision
+accelerate
+spaces
+gradio