Update app.py

app.py

@@ -1,11 +1,11 @@
 import gradio as gr 
-from stability_sdk import client
-import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 from PIL import Image
 import io
 import os
 import warnings
-
+import torch
+import numpy as np
+from diffusers import DiffusionPipeline 
 
 # theme = gr.themes.Monochrome(
 #     primary_hue="indigo",
@@ -15,41 +15,17 @@ import warnings
 #     font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
 # )
 
-def infer(prompt, api_key):
-    stability_api = client.StabilityInference(
-    key=api_key, # API Key reference.
-    verbose=True, # Print debug messages.
-    engine="stable-diffusion-xl-beta-v2-2-2", # Set the engine to use for generation. 
-    # Available engines: stable-diffusion-v1 stable-diffusion-v1-5 stable-diffusion-512-v2-0 stable-diffusion-768-v2-0 stable-inpainting-v1-0 stable-inpainting-512-v2-0
-)
-    answers = stability_api.generate(
-    prompt=prompt,
-    #seed=992446758, # If a seed is provided, the resulting generated image will be deterministic.
-                    # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again.
-                    # Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook.
-    steps=50, # Amount of inference steps performed on image generation. Defaults to 30. 
-    cfg_scale=8.0, # Influences how strongly your generation is guided to match your prompt.
-                   # Setting this value higher increases the strength in which it tries to match your prompt.
-                   # Defaults to 7.0 if not specified.
-    width=512, # Generation width, defaults to 512 if not included.
-    height=512, # Generation height, defaults to 512 if not included.
-    samples=2, # Number of images to generate, defaults to 1 if not included.
-    sampler=generation.SAMPLER_K_DPMPP_2M # Choose which sampler we want to denoise our generation with.
-                                                 # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers.
-                                                 # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m)
-)
-    results = []
-    for resp in answers:
-        for artifact in resp.artifacts:
-            if artifact.finish_reason == generation.FILTER:
-                warnings.warn(
-                "Your request activated the API's safety filters and could not be processed."
-                "Please modify the prompt and try again.")
-            if artifact.type == generation.ARTIFACT_IMAGE:
-                img = Image.open(io.BytesIO(artifact.binary))
-                results.append(img)
+def infer(prompt):
+    torch.cuda.max_memory_allocated(device='cpu')
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-0.9", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    pipe = pipe.to(device)
+    pipe.enable_xformers_memory_efficient_attention()
+    torch.cuda.empty_cache()
+    generator = torch.Generator(device=device).manual_seed(seed)
+    int_image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=steps, height=height, width=width, guidance_scale=scale, num_images_per_prompt=1, generator=generator).images
+    torch.cuda.empty_cache()
     
-    return results
+    return int_image
 
 css = """
         .gradio-container {
@@ -238,8 +214,7 @@ with gr.Blocks(css = css) as demo:
     )
     gr.HTML("<p>You can duplicate this Space to run it privately without a queue for shorter queue times  : <a style='display:inline-block' href='https://huggingface.co/spaces/RamAnanth1/stable-diffusion-xl?duplicate=true'><img src='https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14' alt='Duplicate Space'></a> </p>")
 
-    api_key_input = gr.Textbox(type = "password", label = "Enter your StabilityAI API key here")
-    
+
     with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
         with gr.Column():
             text = gr.Textbox(
@@ -263,7 +238,7 @@ with gr.Blocks(css = css) as demo:
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
-    btn.click(infer, inputs=[text, api_key_input], outputs=[gallery])
+    btn.click(infer, inputs=[text], outputs=[gallery])
     examples = [
         ["Vintage hot rod with custom flame paint job"],
         ["Ancient, mysterious temple in a mountain range, surrounded by misty clouds and tall peaks"],