chore: creating app

app.py

@@ -1,61 +1,58 @@
+import gradio as gr
 import torch
-import spaces
-from PIL import Image
 from diffusers import DiffusionPipeline, AutoencoderKL
-from gradio import Interface
+from PIL import Image
+import spaces  # Assuming you're still working within a framework that uses @spaces
 
-# Load models (outside of the app function for efficiency)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix")
+# Initialize the VAE model and Diffusion Pipeline outside the GPU-enabled function for efficiency
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipe = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     vae=vae,
     torch_dtype=torch.float16,
     variant="fp16",
+    use_safetensors=True
 )
-pipe.load_lora_weights("ritwikraha/khabib_sketch_LoRA")  # Assuming correct model ID
-
-# Move models to CUDA if available (outside of the app function for efficiency)
+pipe.load_lora_weights('ritwikraha/khabib_sketch_LoRA')
 if torch.cuda.is_available():
-    pipe.to("cuda")
+    _ = pipe.to("cuda")
 
-@spaces.GPU
-def inference(prompt, negative_prompt=None, guidance_scale=3, num_inference_steps=50):
-    """Generates an image using the Stable Diffusion XL model with LoRA weights.
+# Define the image generation function
+@spaces.GPU(enable_queue=True)
+def generate_sketch(prompt, negative_prompt="ugly face, multiple bodies, bad anatomy, disfigured, extra fingers", guidance_scale=3, num_inference_steps=50):
+    """Generate a sketch image based on a prompt using Stable Diffusion XL with LoRA weights.
 
     Args:
-        prompt (str): Prompt for image generation, entered by the user.
-        negative_prompt (str, optional): Negative prompt to guide model away from unwanted features. Defaults to "ugly face, multiple bodies, bad anatomy, disfigured, extra fingers".
-        guidance_scale (float, optional): Controls the strength of the guidance from the prompt. Defaults to 3.
-        num_inference_steps (int, optional): Number of inference steps for image generation. Defaults to 50.
+        prompt (str): Description of the image to generate.
+        negative_prompt (str, optional): Negative prompt to avoid certain features. Defaults to common undesirables.
+        guidance_scale (int, optional): The strength of the guidance. Defaults to 3.
+        num_inference_steps (int, optional): The number of steps for the diffusion process. Defaults to 50.
 
     Returns:
-        PIL.Image: Generated image.
+        PIL.Image: The generated sketch image.
     """
-
-    image = pipe(
+    result = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-    ).images[0]
-
-    return image.convert("RGB")  # Ensure RGB format for compatibility
-
-
-
-# Create the Gradio interface
-interface = Interface(
-    fn=inference,
-    inputs=[
-        "text",  # Prompt from user
-        "text",  # Optional negative prompt
-        {"type": "slider", "min": 1, "max": 10, "default": 3},
-        {"type": "slider", "min": 10, "max": 100, "default": 50},
-    ],
-    outputs="image",
-    title="Stable Diffusion XL with Khabib LoRA",
-    description="Generate sketches using the Stable Diffusion XL model fine-tuned on Khabib Nurmagomedov sketches.",
-)
-
-# Launch the Space
-interface.launch()
+    )
+    return result.images[0].convert("RGB")  # Ensure the image is in RGB format
+
+# Setup Gradio interface
+with gr.Blocks() as demo:
+    with gr.Column():
+        prompt_input = gr.Textbox(label="Enter your image prompt", value="a sketch of TOK khabib pointing at another khabib like the spiderman meme, monchrome, pen sketch")
+        negative_prompt_input = gr.Textbox(label="Enter negative prompt", value="ugly face, multiple bodies, bad anatomy, disfigured, extra fingers", lines=2)
+        guidance_scale_slider = gr.Slider(label="Guidance Scale", min_value=1, max_value=5, value=3)
+        steps_slider = gr.Slider(label="Number of Inference Steps", min_value=20, max_value=100, value=50)
+        submit_button = gr.Button("Generate Sketch")
+        output_image = gr.Image(label="Generated Sketch")
+    
+    submit_button.click(
+        fn=generate_sketch,
+        inputs=[prompt_input, negative_prompt_input, guidance_scale_slider, steps_slider],
+        outputs=output_image
+    )
+
+demo.launch()