Add application file

README.md

@@ -1,13 +1,13 @@
 ---
 title: IllusionDiffusion
 emoji: 🔥
-colorFrom: red
+colorFrom: green
 colorTo: pink
 sdk: gradio
 sdk_version: 4.36.1
 app_file: app.py
 pinned: false
-license: openrail
+license: mit
 hf_oauth: true
 disable_embedding: true
 short_description: Generate stunning high quality illusion artwork

app.py

@@ -28,15 +28,15 @@ from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionS
 BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"
 
 # Initialize both pipelines
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
-controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float32)
+controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float32)
 
 # Initialize the safety checker conditionally
 SAFETY_CHECKER_ENABLED = os.environ.get("SAFETY_CHECKER", "0") == "1"
 safety_checker = None
 feature_extractor = None
 if SAFETY_CHECKER_ENABLED:
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to("cuda")
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to("cpu")
     feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
@@ -45,33 +45,11 @@ main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
     vae=vae,
     safety_checker=safety_checker,
     feature_extractor=feature_extractor,
-    torch_dtype=torch.float16,
-).to("cuda")
+    torch_dtype=torch.float32,
+).to("cpu")
 
-# Function to check NSFW images
-#def check_nsfw_images(images: list[Image.Image]) -> tuple[list[Image.Image], list[bool]]:
-#    if SAFETY_CHECKER_ENABLED:
-#        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
-#        has_nsfw_concepts = safety_checker(
-#            images=[images],
-#            clip_input=safety_checker_input.pixel_values.to("cuda")
-#        )
-#        return images, has_nsfw_concepts
-#    else:
-#        return images, [False] * len(images)
-        
-#main_pipe.unet = torch.compile(main_pipe.unet, mode="reduce-overhead", fullgraph=True)
-#main_pipe.unet.to(memory_format=torch.channels_last)
-#main_pipe.unet = torch.compile(main_pipe.unet, mode="reduce-overhead", fullgraph=True)
-#model_id = "stabilityai/sd-x2-latent-upscaler"
 image_pipe = StableDiffusionControlNetImg2ImgPipeline(**main_pipe.components)
 
-
-#image_pipe.unet = torch.compile(image_pipe.unet, mode="reduce-overhead", fullgraph=True)
-#upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-#upscaler.to("cuda")
-
-
 # Sampler map
 SAMPLER_MAP = {
     "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
@@ -113,7 +91,6 @@ def common_upscale(samples, width, height, upscale_method, crop=False):
         return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
 
 def upscale(samples, upscale_method, scale_by):
-        #s = samples.copy()
         width = round(samples["images"].shape[3] * scale_by)
         height = round(samples["images"].shape[2] * scale_by)
         s = common_upscale(samples["images"], width, height, upscale_method, "disabled")
@@ -135,7 +112,6 @@ def convert_to_base64(pil_image):
     return temp_file.name
 
 # Inference function
-@spaces.GPU
 def inference(
     control_image: Image.Image,
     prompt: str,
@@ -155,16 +131,12 @@ def inference(
     start_time_formatted = time.strftime("%H:%M:%S", start_time_struct)
     print(f"Inference started at {start_time_formatted}")
     
-    # Generate the initial image
-    #init_image = init_pipe(prompt).images[0]
-
-    # Rest of your existing code
     control_image_small = center_crop_resize(control_image)
     control_image_large = center_crop_resize(control_image, (1024, 1024))
 
     main_pipe.scheduler = SAMPLER_MAP[sampler](main_pipe.scheduler.config)
     my_seed = random.randint(0, 2**32 - 1) if seed == -1 else seed
-    generator = torch.Generator(device="cuda").manual_seed(my_seed)
+    generator = torch.Generator(device="cpu").manual_seed(my_seed)
     
     out = main_pipe(
         prompt=prompt,
@@ -230,7 +202,6 @@ with gr.Blocks() as app:
         '''
     )
 
-
     state_img_input = gr.State()
     state_img_output = gr.State()
     with gr.Row():
@@ -282,7 +253,7 @@ with gr.Blocks(css=css) as app_with_history:
     with gr.Tab("Past generations"):
         user_history.render()
 
-app_with_history.queue(max_size=20,api_open=False )
+app_with_history.queue(max_size=20,api_open=False)
 
 if __name__ == "__main__":
     app_with_history.launch(max_threads=400,share=True)